From 4d02fa279b23918d8883e57de3c972b3de15c668 Mon Sep 17 00:00:00 2001 From: Hamish Peebles Date: Wed, 24 Jul 2024 12:06:00 +0100 Subject: [PATCH] Store CHIT balances per month (#6087) --- backend/canisters/user/CHANGELOG.md | 1 + backend/canisters/user/api/can.did | 3 + .../user/api/src/queries/chit_events.rs | 1 + .../user/api/src/queries/initial_state.rs | 1 + .../canisters/user/api/src/queries/updates.rs | 1 + backend/canisters/user/impl/src/lib.rs | 7 +- .../user/impl/src/lifecycle/post_upgrade.rs | 2 +- backend/canisters/user/impl/src/model/chit.rs | 168 ++++++++++-------- .../user/impl/src/queries/chit_events.rs | 8 +- .../user/impl/src/queries/initial_state.rs | 3 +- .../user/impl/src/queries/updates.rs | 9 +- .../user/impl/src/updates/claim_daily_chit.rs | 6 +- backend/canisters/user_index/CHANGELOG.md | 1 + backend/canisters/user_index/api/can.did | 15 ++ backend/canisters/user_index/api/src/main.rs | 1 + .../api/src/queries/chit_balances.rs | 21 +++ .../user_index/api/src/queries/mod.rs | 1 + backend/canisters/user_index/impl/src/lib.rs | 15 +- .../impl/src/lifecycle/post_upgrade.rs | 4 +- .../user_index/impl/src/model/user.rs | 33 +++- .../user_index/impl/src/model/user_map.rs | 33 +++- .../src/model/user_referral_leaderboards.rs | 32 +--- .../impl/src/queries/chit_balances.rs | 23 +++ .../impl/src/queries/http_request.rs | 5 +- .../user_index/impl/src/queries/mod.rs | 1 + .../impl/src/queries/referral_leaderboard.rs | 2 +- .../user_index/impl/src/queries/users.rs | 6 +- backend/integration_tests/src/chit_tests.rs | 59 +++++- backend/integration_tests/src/client/user.rs | 1 + .../src/client/user_index.rs | 39 ++++ backend/libraries/types/can.did | 2 + backend/libraries/types/src/user_summary.rs | 2 + backend/libraries/utils/src/time.rs | 91 +++++++++- 33 files changed, 454 insertions(+), 143 deletions(-) create mode 100644 backend/canisters/user_index/api/src/queries/chit_balances.rs create mode 100644 backend/canisters/user_index/impl/src/queries/chit_balances.rs diff --git a/backend/canisters/user/CHANGELOG.md b/backend/canisters/user/CHANGELOG.md index 6dfdc7b890..bc6d373e2e 100644 --- a/backend/canisters/user/CHANGELOG.md +++ b/backend/canisters/user/CHANGELOG.md @@ -13,6 +13,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/). ### Changed - Improve check for empty and dormant users ([#6073](https://github.com/open-chat-labs/open-chat/pull/6073)) +- Store CHIT balances per month ([#6087](https://github.com/open-chat-labs/open-chat/pull/6087)) ## [[2.0.1243](https://github.com/open-chat-labs/open-chat/releases/tag/v2.0.1243-user)] - 2024-07-17 diff --git a/backend/canisters/user/api/can.did b/backend/canisters/user/api/can.did index 8c266f37e6..f95785f921 100644 --- a/backend/canisters/user/api/can.did +++ b/backend/canisters/user/api/can.did @@ -879,6 +879,7 @@ type InitialStateResponse = variant { local_user_index_canister_id : CanisterId; achievements : vec ChitEarned; achievements_last_seen : TimestampMillis; + total_chit_earned : int32; chit_balance : int32; streak : nat16; streak_ends : TimestampMillis; @@ -940,6 +941,7 @@ type UpdatesResponse = variant { suspended : opt bool; achievements : vec ChitEarned; achievements_last_seen : opt TimestampMillis; + total_chit_earned : int32; chit_balance : int32; streak : nat16; streak_ends : TimestampMillis; @@ -1152,6 +1154,7 @@ type LocalUserIndexResponse = variant { type ChitEventsArgs = record { from : opt TimestampMillis; to : opt TimestampMillis; + skip : opt nat32; max : nat32; ascending : bool; }; diff --git a/backend/canisters/user/api/src/queries/chit_events.rs b/backend/canisters/user/api/src/queries/chit_events.rs index f4adf505c9..5ad3cecf64 100644 --- a/backend/canisters/user/api/src/queries/chit_events.rs +++ b/backend/canisters/user/api/src/queries/chit_events.rs @@ -6,6 +6,7 @@ use types::{ChitEarned, TimestampMillis}; pub struct Args { pub from: Option, pub to: Option, + pub skip: Option, pub max: u32, pub ascending: bool, } diff --git a/backend/canisters/user/api/src/queries/initial_state.rs b/backend/canisters/user/api/src/queries/initial_state.rs index 16a54dcaea..5b466f785f 100644 --- a/backend/canisters/user/api/src/queries/initial_state.rs +++ b/backend/canisters/user/api/src/queries/initial_state.rs @@ -23,6 +23,7 @@ pub struct SuccessResult { pub local_user_index_canister_id: CanisterId, pub achievements: Vec, pub achievements_last_seen: TimestampMillis, + pub total_chit_earned: i32, pub chit_balance: i32, pub streak: u16, pub streak_ends: TimestampMillis, diff --git a/backend/canisters/user/api/src/queries/updates.rs b/backend/canisters/user/api/src/queries/updates.rs index f67655904a..dcb5d52884 100644 --- a/backend/canisters/user/api/src/queries/updates.rs +++ b/backend/canisters/user/api/src/queries/updates.rs @@ -32,6 +32,7 @@ pub struct SuccessResult { pub pin_number_settings: OptionUpdate, pub achievements: Vec, pub achievements_last_seen: Option, + pub total_chit_earned: i32, pub chit_balance: i32, pub streak: u16, pub streak_ends: TimestampMillis, diff --git a/backend/canisters/user/impl/src/lib.rs b/backend/canisters/user/impl/src/lib.rs index b33ad87ddb..04bc587dd7 100644 --- a/backend/canisters/user/impl/src/lib.rs +++ b/backend/canisters/user/impl/src/lib.rs @@ -173,7 +173,7 @@ impl RuntimeState { escrow: self.data.escrow_canister_id, icp_ledger: Cryptocurrency::InternetComputer.ledger_canister_id().unwrap(), }, - chit_balance: self.data.chit_balance.value, + chit_balance: self.data.chit_events.balance_for_month_by_timestamp(now), streak: self.data.streak.days(now), streak_ends: self.data.streak.ends(), next_daily_claim: if self.data.streak.can_claim(now) { today(now) } else { tomorrow(now) }, @@ -222,7 +222,6 @@ struct Data { pub pin_number: PinNumber, pub btc_address: Option, pub chit_events: ChitEarnedEvents, - pub chit_balance: Timestamped, pub streak: Streak, pub achievements: HashSet, pub achievements_last_seen: TimestampMillis, @@ -285,7 +284,6 @@ impl Data { pin_number: PinNumber::default(), btc_address: None, chit_events: ChitEarnedEvents::default(), - chit_balance: Timestamped::default(), streak: Streak::default(), achievements: HashSet::new(), achievements_last_seen: 0, @@ -360,7 +358,6 @@ impl Data { timestamp: now, reason: ChitEarnedReason::Achievement(achievement), }); - self.chit_balance = Timestamped::new(self.chit_balance.value + amount, now); true } else { false @@ -370,7 +367,7 @@ impl Data { pub fn notify_user_index_of_chit(&self, now: TimestampMillis) { let args = user_index_canister::c2c_notify_chit::Args { timestamp: now, - chit_balance: self.chit_balance.value, + chit_balance: self.chit_events.balance_for_month_by_timestamp(now), streak: self.streak.days(now), streak_ends: self.streak.ends(), }; diff --git a/backend/canisters/user/impl/src/lifecycle/post_upgrade.rs b/backend/canisters/user/impl/src/lifecycle/post_upgrade.rs index f48bf922da..e83adaba65 100644 --- a/backend/canisters/user/impl/src/lifecycle/post_upgrade.rs +++ b/backend/canisters/user/impl/src/lifecycle/post_upgrade.rs @@ -36,7 +36,7 @@ fn post_upgrade(args: Args) { && state.data.unique_person_proof.is_none() { let now = state.env.now(); - if state.data.user_created + SIX_MONTHS < now && state.data.chit_balance.timestamp + SIX_MONTHS < now { + if state.data.user_created + SIX_MONTHS < now && state.data.chit_events.last_updated() + SIX_MONTHS < now { ic_cdk_timers::set_timer(Duration::ZERO, mark_user_canister_empty); } } diff --git a/backend/canisters/user/impl/src/model/chit.rs b/backend/canisters/user/impl/src/model/chit.rs index 0dcf4a8b03..794e2da003 100644 --- a/backend/canisters/user/impl/src/model/chit.rs +++ b/backend/canisters/user/impl/src/model/chit.rs @@ -1,9 +1,18 @@ use serde::{Deserialize, Serialize}; +use std::ops::Range; use types::{ChitEarned, ChitEarnedReason, TimestampMillis}; +use utils::time::MonthKey; #[derive(Serialize, Deserialize, Default)] +#[serde(from = "ChitEarnedEventsPrevious")] pub struct ChitEarnedEvents { events: Vec, + total_chit_earned: i32, +} + +#[derive(Serialize, Deserialize, Default)] +pub struct ChitEarnedEventsPrevious { + events: Vec, } impl ChitEarnedEvents { @@ -16,6 +25,7 @@ impl ChitEarnedEvents { } } + self.total_chit_earned += event.amount; self.events.push(event); if sort { @@ -27,29 +37,31 @@ impl ChitEarnedEvents { &self, from: Option, to: Option, - max: u32, + skip: usize, + max: usize, ascending: bool, ) -> (Vec, u32) { - let page = if ascending { - self.events - .iter() - .skip_while(|e| from.map_or(false, |ts| e.timestamp <= ts)) - .take_while(|e| to.map_or(true, |ts| e.timestamp <= ts)) - .take(max as usize) - .cloned() - .collect() + if ascending { + let range = self.range(from.unwrap_or_default()..to.unwrap_or(TimestampMillis::MAX)); + (range.iter().skip(skip).take(max).cloned().collect(), range.len() as u32) } else { - self.events - .iter() - .rev() - .skip_while(|e| from.map_or(false, |ts| e.timestamp >= ts)) - .take_while(|e| to.map_or(true, |ts| e.timestamp >= ts)) - .take(max as usize) - .cloned() - .collect() - }; + let range = self.range(to.unwrap_or_default()..from.unwrap_or(TimestampMillis::MAX)); + (range.iter().rev().skip(skip).take(max).cloned().collect(), range.len() as u32) + } + } - (page, self.events.len() as u32) + pub fn total_chit_earned(&self) -> i32 { + self.total_chit_earned + } + + pub fn balance_for_month_by_timestamp(&self, ts: TimestampMillis) -> i32 { + self.balance_for_month(MonthKey::from_timestamp(ts)) + } + + pub fn balance_for_month(&self, month: MonthKey) -> i32 { + let timestamp_range = month.timestamp_range(); + let range = self.range(timestamp_range); + range.iter().map(|e| e.amount).sum() } pub fn achievements(&self, since: Option) -> Vec { @@ -62,12 +74,26 @@ impl ChitEarnedEvents { .collect() } - pub fn has_achievements_since(&self, since: TimestampMillis) -> bool { - self.events - .iter() - .rev() - .take_while(|e| e.timestamp > since) - .any(|e| matches!(e.reason, ChitEarnedReason::Achievement(_))) + pub fn last_updated(&self) -> TimestampMillis { + self.events.last().map(|e| e.timestamp).unwrap_or_default() + } + + fn range(&self, range: Range) -> &[ChitEarned] { + let start = self.events.partition_point(|e| e.timestamp < range.start); + let end = self.events.partition_point(|e| e.timestamp <= range.end); + + &self.events[start..end] + } +} + +impl From for ChitEarnedEvents { + fn from(value: ChitEarnedEventsPrevious) -> Self { + let total_chit_earned = value.events.iter().map(|e| e.amount).sum(); + + ChitEarnedEvents { + events: value.events, + total_chit_earned, + } } } @@ -81,7 +107,7 @@ mod tests { fn first_page_matches_expected() { let store = init_test_data(); - let (events, total) = store.events(None, None, 3, true); + let (events, total) = store.events(None, None, 0, 3, true); assert_eq!(total, 7); assert_eq!(events.len(), 3); @@ -93,8 +119,7 @@ mod tests { fn next_page_matches_expected() { let store = init_test_data(); - let (events, _) = store.events(None, None, 3, true); - let (events, _) = store.events(Some(events[2].timestamp), None, 3, true); + let (events, _) = store.events(None, None, 3, 3, true); assert_eq!(events.len(), 3); assert_eq!(events[0].timestamp, 13); @@ -105,7 +130,7 @@ mod tests { fn first_page_desc_matches_expected() { let store = init_test_data(); - let (events, _) = store.events(None, None, 3, false); + let (events, _) = store.events(None, None, 0, 3, false); assert_eq!(events.len(), 3); assert_eq!(events[0].timestamp, 16); @@ -116,8 +141,7 @@ mod tests { fn next_page_desc_matches_expected() { let store = init_test_data(); - let (events, _) = store.events(None, None, 3, false); - let (events, _) = store.events(Some(events[2].timestamp), None, 3, false); + let (events, _) = store.events(None, None, 3, 3, false); assert_eq!(events.len(), 3); assert_eq!(events[0].timestamp, 13); @@ -128,7 +152,7 @@ mod tests { fn range_matches_expected() { let store = init_test_data(); - let (events, _) = store.events(Some(11), Some(15), 99, true); + let (events, _) = store.events(Some(12), Some(15), 0, 99, true); assert_eq!(events.len(), 4); assert_eq!(events[0].timestamp, 12); @@ -139,7 +163,7 @@ mod tests { fn range_desc_matches_expected() { let store = init_test_data(); - let (events, _) = store.events(Some(15), Some(11), 99, false); + let (events, _) = store.events(Some(14), Some(11), 0, 99, false); assert_eq!(events.len(), 4); assert_eq!(events[0].timestamp, 14); @@ -147,44 +171,48 @@ mod tests { } fn init_test_data() -> ChitEarnedEvents { + let events = vec![ + ChitEarned { + amount: 200, + timestamp: 10, + reason: ChitEarnedReason::DailyClaim, + }, + ChitEarned { + amount: 200, + timestamp: 11, + reason: ChitEarnedReason::DailyClaim, + }, + ChitEarned { + amount: 300, + timestamp: 12, + reason: ChitEarnedReason::DailyClaim, + }, + ChitEarned { + amount: 500, + timestamp: 13, + reason: ChitEarnedReason::Achievement(Achievement::SetBio), + }, + ChitEarned { + amount: 300, + timestamp: 14, + reason: ChitEarnedReason::DailyClaim, + }, + ChitEarned { + amount: 500, + timestamp: 15, + reason: ChitEarnedReason::Achievement(Achievement::SetAvatar), + }, + ChitEarned { + amount: 500, + timestamp: 16, + reason: ChitEarnedReason::Achievement(Achievement::SentDirectMessage), + }, + ]; + let total_chit_earned = events.iter().map(|e| e.amount).sum(); + ChitEarnedEvents { - events: vec![ - ChitEarned { - amount: 200, - timestamp: 10, - reason: ChitEarnedReason::DailyClaim, - }, - ChitEarned { - amount: 200, - timestamp: 11, - reason: ChitEarnedReason::DailyClaim, - }, - ChitEarned { - amount: 300, - timestamp: 12, - reason: ChitEarnedReason::DailyClaim, - }, - ChitEarned { - amount: 500, - timestamp: 13, - reason: ChitEarnedReason::Achievement(Achievement::SetBio), - }, - ChitEarned { - amount: 300, - timestamp: 14, - reason: ChitEarnedReason::DailyClaim, - }, - ChitEarned { - amount: 500, - timestamp: 15, - reason: ChitEarnedReason::Achievement(Achievement::SetAvatar), - }, - ChitEarned { - amount: 500, - timestamp: 16, - reason: ChitEarnedReason::Achievement(Achievement::SentDirectMessage), - }, - ], + events, + total_chit_earned, } } } diff --git a/backend/canisters/user/impl/src/queries/chit_events.rs b/backend/canisters/user/impl/src/queries/chit_events.rs index d5b7652493..86ccf74902 100644 --- a/backend/canisters/user/impl/src/queries/chit_events.rs +++ b/backend/canisters/user/impl/src/queries/chit_events.rs @@ -9,7 +9,13 @@ fn chit_events(args: Args) -> Response { } fn chit_events_impl(args: Args, state: &RuntimeState) -> Response { - let (events, total) = state.data.chit_events.events(args.from, args.to, args.max, args.ascending); + let (events, total) = state.data.chit_events.events( + args.from, + args.to, + args.skip.unwrap_or_default() as usize, + args.max as usize, + args.ascending, + ); Response::Success(SuccessResult { events, total }) } diff --git a/backend/canisters/user/impl/src/queries/initial_state.rs b/backend/canisters/user/impl/src/queries/initial_state.rs index d7a47779d0..5317274342 100644 --- a/backend/canisters/user/impl/src/queries/initial_state.rs +++ b/backend/canisters/user/impl/src/queries/initial_state.rs @@ -49,7 +49,8 @@ fn initial_state_impl(state: &RuntimeState) -> Response { local_user_index_canister_id: state.data.local_user_index_canister_id, achievements: state.data.chit_events.achievements(None), achievements_last_seen: state.data.achievements_last_seen, - chit_balance: state.data.chit_balance.value, + total_chit_earned: state.data.chit_events.total_chit_earned(), + chit_balance: state.data.chit_events.balance_for_month_by_timestamp(now), streak: state.data.streak.days(now), streak_ends: state.data.streak.ends(), next_daily_claim: if state.data.streak.can_claim(now) { today(now) } else { tomorrow(now) }, diff --git a/backend/canisters/user/impl/src/queries/updates.rs b/backend/canisters/user/impl/src/queries/updates.rs index 3da85d316d..01b0f35dac 100644 --- a/backend/canisters/user/impl/src/queries/updates.rs +++ b/backend/canisters/user/impl/src/queries/updates.rs @@ -53,9 +53,8 @@ fn updates_impl(updates_since: TimestampMillis, state: &RuntimeState) -> Respons || state.data.group_chats.any_updated(updates_since) || state.data.favourite_chats.any_updated(updates_since) || state.data.communities.any_updated(updates_since) - || state.data.chit_events.has_achievements_since(updates_since) - || state.data.achievements_last_seen > updates_since - || state.data.chit_balance.timestamp > updates_since; + || state.data.chit_events.last_updated() > updates_since + || state.data.achievements_last_seen > updates_since; // Short circuit prior to calling `ic0.time()` so that caching works effectively if !has_any_updates { @@ -140,7 +139,8 @@ fn updates_impl(updates_since: TimestampMillis, state: &RuntimeState) -> Respons None }; - let chit_balance = state.data.chit_balance.value; + let total_chit_earned = state.data.chit_events.total_chit_earned(); + let chit_balance = state.data.chit_events.balance_for_month_by_timestamp(now); let streak = state.data.streak.days(now); let next_daily_claim = if state.data.streak.can_claim(now) { today(now) } else { tomorrow(now) }; let streak_ends = state.data.streak.ends(); @@ -160,6 +160,7 @@ fn updates_impl(updates_since: TimestampMillis, state: &RuntimeState) -> Respons pin_number_settings, achievements, achievements_last_seen, + total_chit_earned, chit_balance, streak, streak_ends, diff --git a/backend/canisters/user/impl/src/updates/claim_daily_chit.rs b/backend/canisters/user/impl/src/updates/claim_daily_chit.rs index 15a30910b0..fc27e132d0 100644 --- a/backend/canisters/user/impl/src/updates/claim_daily_chit.rs +++ b/backend/canisters/user/impl/src/updates/claim_daily_chit.rs @@ -4,7 +4,7 @@ use canister_tracing_macros::trace; use event_store_producer::EventBuilder; use ic_cdk::update; use serde::Serialize; -use types::{Achievement, ChitEarned, ChitEarnedReason, Timestamped, UserId}; +use types::{Achievement, ChitEarned, ChitEarnedReason, UserId}; use user_canister::claim_daily_chit::{Response::*, *}; use utils::time::tomorrow; @@ -26,8 +26,6 @@ fn claim_daily_chit_impl(state: &mut RuntimeState) -> Response { let streak = state.data.streak.days(now); let chit_earned = chit_for_streak(streak); - state.data.chit_balance = Timestamped::new(state.data.chit_balance.value + chit_earned as i32, now); - state.data.chit_events.push(ChitEarned { amount: chit_earned as i32, timestamp: now, @@ -62,7 +60,7 @@ fn claim_daily_chit_impl(state: &mut RuntimeState) -> Response { Success(SuccessResult { chit_earned, - chit_balance: state.data.chit_balance.value, + chit_balance: state.data.chit_events.balance_for_month_by_timestamp(now), streak, next_claim: tomorrow, }) diff --git a/backend/canisters/user_index/CHANGELOG.md b/backend/canisters/user_index/CHANGELOG.md index 873f5dcf8e..00849f6f95 100644 --- a/backend/canisters/user_index/CHANGELOG.md +++ b/backend/canisters/user_index/CHANGELOG.md @@ -18,6 +18,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/). - Handle transfer fee changing in either direction ([#6064](https://github.com/open-chat-labs/open-chat/pull/6064)) - Accept proofs of uniqueness from LocalUserIndexes ([#6068](https://github.com/open-chat-labs/open-chat/pull/6068)) - Ensure UserIndex is only controller before installing LocalUserIndex ([#6070](https://github.com/open-chat-labs/open-chat/pull/6070)) +- Store CHIT balances per month ([#6087](https://github.com/open-chat-labs/open-chat/pull/6087)) - Add `user_ii_principal` to `submit_proof_of_unique_personhood` args ([#6092](https://github.com/open-chat-labs/open-chat/pull/6092)) - Bump `ic-verifiable-credentials` to latest version ([#6096](https://github.com/open-chat-labs/open-chat/pull/6096)) diff --git a/backend/canisters/user_index/api/can.did b/backend/canisters/user_index/api/can.did index c74a9ffc6d..466462d1cf 100644 --- a/backend/canisters/user_index/api/can.did +++ b/backend/canisters/user_index/api/can.did @@ -117,6 +117,18 @@ type SearchResponse = variant { }; }; +type ChitBalancesArgs = record { + users : vec UserId; + year : nat16; + month : nat8; +}; + +type ChitBalancesResponse = variant { + Success : record { + balances : vec record { UserId; int32 }; + }; +}; + type DiamondMembershipFeesResponse = variant { Success : vec record { token : Cryptocurrency; @@ -405,6 +417,9 @@ service : { // Search for users matching some query search : (SearchArgs) -> (SearchResponse) query; + // Gets the CHIT balances for multiple users for a chosen month + chit_balances : (ChitBalancesArgs) -> (ChitBalancesResponse) query; + // Retrieves the current fees to pay for Diamond membership diamond_membership_fees : (EmptyArgs) -> (DiamondMembershipFeesResponse) query; diff --git a/backend/canisters/user_index/api/src/main.rs b/backend/canisters/user_index/api/src/main.rs index 1a73eec013..690ad83551 100644 --- a/backend/canisters/user_index/api/src/main.rs +++ b/backend/canisters/user_index/api/src/main.rs @@ -3,6 +3,7 @@ use candid_gen::generate_candid_method; #[allow(deprecated)] fn main() { generate_candid_method!(user_index, check_username, query); + generate_candid_method!(user_index, chit_balances, query); generate_candid_method!(user_index, chit_leaderboard, query); generate_candid_method!(user_index, current_user, query); generate_candid_method!(user_index, diamond_membership_fees, query); diff --git a/backend/canisters/user_index/api/src/queries/chit_balances.rs b/backend/canisters/user_index/api/src/queries/chit_balances.rs new file mode 100644 index 0000000000..d618571cbd --- /dev/null +++ b/backend/canisters/user_index/api/src/queries/chit_balances.rs @@ -0,0 +1,21 @@ +use candid::CandidType; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use types::UserId; + +#[derive(CandidType, Serialize, Deserialize, Debug)] +pub struct Args { + pub users: Vec, + pub year: u16, + pub month: u8, +} + +#[derive(CandidType, Serialize, Deserialize, Debug)] +pub enum Response { + Success(SuccessResult), +} + +#[derive(CandidType, Serialize, Deserialize, Debug)] +pub struct SuccessResult { + pub balances: HashMap, +} diff --git a/backend/canisters/user_index/api/src/queries/mod.rs b/backend/canisters/user_index/api/src/queries/mod.rs index a87ddbb581..5b344c9bdd 100644 --- a/backend/canisters/user_index/api/src/queries/mod.rs +++ b/backend/canisters/user_index/api/src/queries/mod.rs @@ -1,5 +1,6 @@ pub mod c2c_lookup_user; pub mod check_username; +pub mod chit_balances; pub mod chit_leaderboard; pub mod current_user; pub mod diamond_membership_fees; diff --git a/backend/canisters/user_index/impl/src/lib.rs b/backend/canisters/user_index/impl/src/lib.rs index 1f74259c69..9bce3d45cb 100644 --- a/backend/canisters/user_index/impl/src/lib.rs +++ b/backend/canisters/user_index/impl/src/lib.rs @@ -33,7 +33,7 @@ use utils::canister::{CanistersRequiringUpgrade, FailedUpgradeCount}; use utils::canister_event_sync_queue::CanisterEventSyncQueue; use utils::consts::DEV_TEAM_DFX_PRINCIPAL; use utils::env::Environment; -use utils::time::DAY_IN_MS; +use utils::time::{MonthKey, DAY_IN_MS}; mod guards; mod jobs; @@ -199,7 +199,11 @@ impl RuntimeState { is_bot: user.is_bot, suspension_details: user.suspension_details.clone(), moderation_flags_enabled: user.moderation_flags_enabled, - chit_balance: user.chit_balance, + chit_balance: user + .chit_per_month + .get(&MonthKey::from_timestamp(now)) + .copied() + .unwrap_or_default(), streak: user.streak(now), streak_ends: user.streak_ends, chit_updated: user.chit_updated, @@ -426,13 +430,16 @@ impl Data { } } - pub fn chit_bands(&self, size: u32) -> BTreeMap { + pub fn chit_bands(&self, size: u32, year: u32, month: u8) -> BTreeMap { let mut bands = BTreeMap::new(); + let month_key = MonthKey::new(year, month); for chit in self .users .iter() - .map(|u| if u.chit_balance > 0 { u.chit_balance as u32 } else { 0 }) + .map(|u| u.chit_per_month.get(&month_key).copied().unwrap_or_default()) + .filter(|c| *c > 0) + .map(|c| c as u32) { let band = (chit / size) * size; let key = if band > 0 { (chit / band) * band } else { 0 }; diff --git a/backend/canisters/user_index/impl/src/lifecycle/post_upgrade.rs b/backend/canisters/user_index/impl/src/lifecycle/post_upgrade.rs index c163dddb67..b4f420cf9e 100644 --- a/backend/canisters/user_index/impl/src/lifecycle/post_upgrade.rs +++ b/backend/canisters/user_index/impl/src/lifecycle/post_upgrade.rs @@ -1,6 +1,6 @@ use crate::lifecycle::{init_env, init_state}; use crate::memory::get_upgrades_memory; -use crate::Data; +use crate::{mutate_state, Data}; use canister_logger::LogEntry; use canister_tracing_macros::trace; use ic_cdk::post_upgrade; @@ -23,5 +23,7 @@ fn post_upgrade(args: Args) { init_cycles_dispenser_client(data.cycles_dispenser_canister_id, data.test_mode); init_state(env, data, args.wasm_version); + mutate_state(|state| state.data.users.initialise_monthly_chit_balances(state.env.now())); + info!(version = %args.wasm_version, "Post-upgrade complete"); } diff --git a/backend/canisters/user_index/impl/src/model/user.rs b/backend/canisters/user_index/impl/src/model/user.rs index d2a3d2aaba..e56731f7ac 100644 --- a/backend/canisters/user_index/impl/src/model/user.rs +++ b/backend/canisters/user_index/impl/src/model/user.rs @@ -2,11 +2,13 @@ use crate::model::diamond_membership_details::DiamondMembershipDetailsInternal; use crate::{model::account_billing::AccountBilling, TIME_UNTIL_SUSPENDED_ACCOUNT_IS_DELETED_MILLIS}; use candid::{CandidType, Principal}; use serde::{Deserialize, Serialize}; +use std::collections::BTreeMap; use types::{ is_default, is_empty_slice, CyclesTopUp, CyclesTopUpInternal, PhoneNumber, RegistrationFee, SuspensionAction, SuspensionDuration, TimestampMillis, UniquePersonProof, UserId, UserSummary, UserSummaryStable, UserSummaryV2, UserSummaryVolatile, }; +use utils::time::MonthKey; #[derive(Serialize, Deserialize, Clone)] pub struct User { @@ -51,7 +53,10 @@ pub struct User { #[serde(rename = "rm", default, skip_serializing_if = "is_empty_slice")] pub reported_messages: Vec, #[serde(rename = "cb", alias = "c2", default, skip_serializing_if = "is_default")] + #[deprecated] pub chit_balance: i32, + #[serde(rename = "cm", default, skip_serializing_if = "is_default")] + pub chit_per_month: BTreeMap, #[serde(rename = "sk", alias = "s", default, skip_serializing_if = "is_default")] pub streak: u16, #[serde(rename = "se", default, skip_serializing_if = "is_default")] @@ -60,6 +65,8 @@ pub struct User { pub chit_updated: TimestampMillis, #[serde(rename = "lc", default)] pub latest_chit_event: TimestampMillis, + #[serde(rename = "lcp", default)] + pub latest_chit_event_previous_month: TimestampMillis, #[serde(rename = "uh", default, skip_serializing_if = "Option::is_none")] pub unique_person_proof: Option, } @@ -73,6 +80,10 @@ impl User { pub fn mark_cycles_top_up(&mut self, top_up: CyclesTopUp) { self.cycle_top_ups.push(top_up.into()) } + + pub fn total_chit_earned(&self) -> i32 { + self.chit_per_month.values().copied().sum() + } } #[derive(CandidType, Serialize, Deserialize, Clone, Debug, Default, Eq, PartialEq)] @@ -92,6 +103,7 @@ impl User { referred_by: Option, is_bot: bool, ) -> User { + #[allow(deprecated)] User { principal, user_id, @@ -112,10 +124,12 @@ impl User { moderation_flags_enabled: 0, reported_messages: Vec::new(), chit_balance: 0, + chit_per_month: BTreeMap::new(), chit_updated: now, streak: 0, streak_ends: 0, latest_chit_event: 0, + latest_chit_event_previous_month: 0, unique_person_proof: None, } } @@ -130,17 +144,22 @@ impl User { suspended: self.suspension_details.is_some(), diamond_member: self.diamond_membership_details.is_active(now), diamond_membership_status: self.diamond_membership_details.status(now), - chit_balance: self.chit_balance, + total_chit_earned: self.total_chit_earned(), + chit_balance: self + .chit_per_month + .get(&MonthKey::from_timestamp(now)) + .copied() + .unwrap_or_default(), streak: self.streak(now), is_unique_person: self.unique_person_proof.is_some(), } } - pub fn to_summary_v2(&self, now: TimestampMillis) -> UserSummaryV2 { + pub fn to_summary_v2(&self, now: TimestampMillis, month_key: MonthKey) -> UserSummaryV2 { UserSummaryV2 { user_id: self.user_id, stable: Some(self.to_summary_stable(now)), - volatile: Some(self.to_summary_volatile(now)), + volatile: Some(self.to_summary_volatile(now, month_key)), } } @@ -156,9 +175,10 @@ impl User { } } - pub fn to_summary_volatile(&self, now: TimestampMillis) -> UserSummaryVolatile { + pub fn to_summary_volatile(&self, now: TimestampMillis, month_key: MonthKey) -> UserSummaryVolatile { UserSummaryVolatile { - chit_balance: self.chit_balance, + total_chit_earned: self.total_chit_earned(), + chit_balance: self.chit_per_month.get(&month_key).copied().unwrap_or_default(), streak: self.streak(now), } } @@ -206,6 +226,7 @@ impl From<&SuspensionDetails> for types::SuspensionDetails { #[cfg(test)] impl Default for User { fn default() -> Self { + #[allow(deprecated)] User { principal: Principal::anonymous(), user_id: Principal::anonymous().into(), @@ -226,10 +247,12 @@ impl Default for User { moderation_flags_enabled: 0, reported_messages: Vec::new(), chit_balance: 0, + chit_per_month: BTreeMap::new(), streak: 0, streak_ends: 0, chit_updated: 0, latest_chit_event: 0, + latest_chit_event_previous_month: 0, unique_person_proof: None, } } diff --git a/backend/canisters/user_index/impl/src/model/user_map.rs b/backend/canisters/user_index/impl/src/model/user_map.rs index 7998b1d210..c45f87f953 100644 --- a/backend/canisters/user_index/impl/src/model/user_map.rs +++ b/backend/canisters/user_index/impl/src/model/user_map.rs @@ -9,6 +9,7 @@ use std::ops::RangeFrom; use tracing::info; use types::{CyclesTopUp, Milliseconds, SuspensionDuration, TimestampMillis, UniquePersonProof, UserId}; use utils::case_insensitive_hash_map::CaseInsensitiveHashMap; +use utils::time::MonthKey; #[derive(Serialize, Deserialize, Default)] #[serde(from = "UserMapTrimmed")] @@ -33,6 +34,14 @@ pub struct UserMap { } impl UserMap { + pub fn initialise_monthly_chit_balances(&mut self, now: TimestampMillis) { + let month_key = MonthKey::from_timestamp(now); + for user in self.users.values_mut() { + #[allow(deprecated)] + user.chit_per_month.insert(month_key, user.chit_balance); + } + } + pub fn does_username_exist(&self, username: &str) -> bool { self.username_to_user_id.contains_key(username) } @@ -201,16 +210,26 @@ impl UserMap { return false; }; - if chit_event_timestamp <= user.latest_chit_event { - return false; + let chit_event_month = MonthKey::from_timestamp(chit_event_timestamp); + + if chit_event_timestamp >= user.latest_chit_event { + if MonthKey::from_timestamp(user.latest_chit_event) == chit_event_month.previous() { + user.latest_chit_event_previous_month = user.latest_chit_event; + } + user.latest_chit_event = chit_event_timestamp; + user.streak = streak; + user.streak_ends = streak_ends; + } else { + let previous_month = MonthKey::from_timestamp(now).previous(); + if chit_event_month == previous_month && chit_event_timestamp >= user.latest_chit_event_previous_month { + user.latest_chit_event_previous_month = chit_event_timestamp; + } else { + return false; + } } - user.latest_chit_event = chit_event_timestamp; - user.chit_balance = chit_balance; - user.streak = streak; - user.streak_ends = streak_ends; user.chit_updated = now; - + user.chit_per_month.insert(chit_event_month, chit_balance); true } diff --git a/backend/canisters/user_index/impl/src/model/user_referral_leaderboards.rs b/backend/canisters/user_index/impl/src/model/user_referral_leaderboards.rs index a418d216a5..310917ee62 100644 --- a/backend/canisters/user_index/impl/src/model/user_referral_leaderboards.rs +++ b/backend/canisters/user_index/impl/src/model/user_referral_leaderboards.rs @@ -4,8 +4,9 @@ use std::collections::binary_heap::BinaryHeap; use std::collections::btree_map::BTreeMap; use std::collections::HashMap; use types::{TimestampMillis, UserId}; +use utils::time::MonthKey; -const STARTING_MONTH: MonthKey = MonthKey { year: 2023, month: 5 }; +const STARTING_MONTH: MonthKey = MonthKey::new(2023, 5); #[derive(Serialize, Deserialize, Default)] pub struct UserReferralLeaderboards { @@ -107,35 +108,6 @@ pub struct ReferralStats { pub user_id: UserId, } -#[derive(Serialize, Deserialize, Copy, Clone, Eq, PartialEq, Ord, PartialOrd)] -pub struct MonthKey { - year: u32, - month: u8, -} - -impl MonthKey { - pub fn new(year: u32, month: u8) -> MonthKey { - MonthKey { year, month } - } - - pub fn from_timestamp(ts: TimestampMillis) -> MonthKey { - let date = time::OffsetDateTime::from_unix_timestamp((ts / 1000) as i64).unwrap(); - - MonthKey { - year: date.year() as u32, - month: u8::from(date.month()), - } - } - - pub fn year(&self) -> u32 { - self.year - } - - pub fn month(&self) -> u8 { - self.month - } -} - #[cfg(test)] mod tests { use super::*; diff --git a/backend/canisters/user_index/impl/src/queries/chit_balances.rs b/backend/canisters/user_index/impl/src/queries/chit_balances.rs new file mode 100644 index 0000000000..0d6d903908 --- /dev/null +++ b/backend/canisters/user_index/impl/src/queries/chit_balances.rs @@ -0,0 +1,23 @@ +use crate::{read_state, RuntimeState}; +use ic_cdk::query; +use user_index_canister::chit_balances::{Response::*, *}; +use utils::time::MonthKey; + +#[query] +fn chit_balances(args: Args) -> Response { + read_state(|state| chit_balances_impl(args, state)) +} + +fn chit_balances_impl(args: Args, state: &RuntimeState) -> Response { + let month_key = MonthKey::new(args.year as u32, args.month); + + let balances = args + .users + .iter() + .flat_map(|u| state.data.users.get_by_user_id(u)) + .map(|u| (u.user_id, u.chit_per_month.get(&month_key).copied().unwrap_or_default())) + .filter(|(_, c)| *c > 0) + .collect(); + + Success(SuccessResult { balances }) +} diff --git a/backend/canisters/user_index/impl/src/queries/http_request.rs b/backend/canisters/user_index/impl/src/queries/http_request.rs index 1e5f8d6b96..37b47ecb2b 100644 --- a/backend/canisters/user_index/impl/src/queries/http_request.rs +++ b/backend/canisters/user_index/impl/src/queries/http_request.rs @@ -4,6 +4,7 @@ use http_request::{build_json_response, encode_logs, extract_route, Route}; use ic_cdk::query; use std::collections::BTreeMap; use types::{HttpRequest, HttpResponse, TimestampMillis, UserId}; +use utils::time::MonthKey; #[query] fn http_request(request: HttpRequest) -> HttpResponse { @@ -57,8 +58,10 @@ fn http_request(request: HttpRequest) -> HttpResponse { "new_users_per_day" => return get_new_users_per_day(state), "chitbands" => { let size: u32 = parts.get(1).and_then(|s| (*s).parse::().ok()).unwrap_or(500); + let now = state.env.now(); + let month_key = MonthKey::from_timestamp(now); - return build_json_response(&state.data.chit_bands(size)); + return build_json_response(&state.data.chit_bands(size, month_key.year(), month_key.month())); } _ => (), } diff --git a/backend/canisters/user_index/impl/src/queries/mod.rs b/backend/canisters/user_index/impl/src/queries/mod.rs index 892ad45472..d9ec2dcac5 100644 --- a/backend/canisters/user_index/impl/src/queries/mod.rs +++ b/backend/canisters/user_index/impl/src/queries/mod.rs @@ -1,5 +1,6 @@ pub mod c2c_lookup_user; pub mod check_username; +pub mod chit_balances; pub mod chit_leaderboard; pub mod current_user; pub mod diamond_membership_fees; diff --git a/backend/canisters/user_index/impl/src/queries/referral_leaderboard.rs b/backend/canisters/user_index/impl/src/queries/referral_leaderboard.rs index b91542f66d..0aeaa783a8 100644 --- a/backend/canisters/user_index/impl/src/queries/referral_leaderboard.rs +++ b/backend/canisters/user_index/impl/src/queries/referral_leaderboard.rs @@ -1,7 +1,7 @@ -use crate::model::user_referral_leaderboards::MonthKey; use crate::{read_state, RuntimeState}; use ic_cdk::query; use user_index_canister::referral_leaderboard::{Response::*, *}; +use utils::time::MonthKey; #[query] fn referral_leaderboard(args: Args) -> Response { diff --git a/backend/canisters/user_index/impl/src/queries/users.rs b/backend/canisters/user_index/impl/src/queries/users.rs index dc19f5ad07..3d81a6dc44 100644 --- a/backend/canisters/user_index/impl/src/queries/users.rs +++ b/backend/canisters/user_index/impl/src/queries/users.rs @@ -3,6 +3,7 @@ use ic_cdk::query; use std::collections::HashSet; use types::{CurrentUserSummary, UserSummaryV2}; use user_index_canister::users::{Response::*, *}; +use utils::time::MonthKey; #[query] fn users(args: Args) -> Response { @@ -46,6 +47,7 @@ fn users_impl(args: Args, state: &RuntimeState) -> Response { } } + let now_month = MonthKey::from_timestamp(now); for group in args.user_groups { let updated_since = group.updated_since; users.extend( @@ -63,7 +65,7 @@ fn users_impl(args: Args, state: &RuntimeState) -> Response { .map(|u| UserSummaryV2 { user_id: u.user_id, stable: (u.date_updated > updated_since).then(|| u.to_summary_stable(now)), - volatile: Some(u.to_summary_volatile(now)), + volatile: Some(u.to_summary_volatile(now, now_month)), }), ); } @@ -78,7 +80,7 @@ fn users_impl(args: Args, state: &RuntimeState) -> Response { .take(100) .filter(|u| user_ids.insert(*u)) .filter_map(|u| state.data.users.get_by_user_id(&u)) - .map(|u| u.to_summary_v2(now)), + .map(|u| u.to_summary_v2(now, now_month)), ); } diff --git a/backend/integration_tests/src/chit_tests.rs b/backend/integration_tests/src/chit_tests.rs index 3f8bdede44..e5b9fcbdc0 100644 --- a/backend/integration_tests/src/chit_tests.rs +++ b/backend/integration_tests/src/chit_tests.rs @@ -1,11 +1,11 @@ -use pocket_ic::PocketIc; -use types::{ChitEarnedReason, Milliseconds, TimestampMillis}; - use crate::env::ENV; use crate::utils::now_millis; use crate::{client, TestEnv}; -use std::ops::Deref; +use pocket_ic::PocketIc; +use std::ops::{Add, Deref}; use std::time::{Duration, SystemTime}; +use types::{ChitEarnedReason, Milliseconds, TimestampMillis}; +use utils::time::MonthKey; const DAY_ZERO: TimestampMillis = 1704067200000; // Mon Jan 01 2024 00:00:00 GMT+0000 const MS_IN_DAY: Milliseconds = 1000 * 60 * 60 * 24; @@ -80,6 +80,57 @@ fn chit_streak_gained_and_lost_as_expected() { assert_eq!(result.users[0].volatile.as_ref().unwrap().streak, 0); } +#[test] +fn chit_stored_per_month() { + let mut wrapper = ENV.deref().get(); + let TestEnv { env, canister_ids, .. } = wrapper.env(); + + let user = client::register_user(env, canister_ids); + ensure_time_at_least_day0(env); + advance_to_next_month(env); + let start_month = MonthKey::from_timestamp(now_millis(env)); + + for i in 1..5 { + for _ in 0..i { + client::user::happy_path::claim_daily_chit(env, &user); + env.advance_time(Duration::from_millis(2 * MS_IN_DAY)); + } + advance_to_next_month(env); + } + + env.tick(); + + let mut month = start_month; + for i in 1..5 { + let chit_balance = client::user_index::happy_path::chit_balances( + env, + canister_ids.user_index, + vec![user.user_id], + month.year() as u16, + month.month(), + ) + .values() + .next() + .copied() + .unwrap(); + + assert_eq!(chit_balance, i * 200); + + month = month.next(); + } + + let user = client::user_index::happy_path::user(env, canister_ids.user_index, user.user_id); + + assert_eq!(user.total_chit_earned, 2000); + assert_eq!(user.chit_balance, 0); +} + +fn advance_to_next_month(env: &mut PocketIc) { + let now = now_millis(env); + let next_month = MonthKey::from_timestamp(now).next(); + env.set_time(SystemTime::UNIX_EPOCH.add(Duration::from_millis(next_month.start_timestamp()))); +} + fn ensure_time_at_least_day0(env: &mut PocketIc) { if now_millis(env) < DAY_ZERO { env.set_time(SystemTime::now()); diff --git a/backend/integration_tests/src/client/user.rs b/backend/integration_tests/src/client/user.rs index f231497bd8..9fba9de9e7 100644 --- a/backend/integration_tests/src/client/user.rs +++ b/backend/integration_tests/src/client/user.rs @@ -403,6 +403,7 @@ pub mod happy_path { &user_canister::chit_events::Args { from, to, + skip: None, max, ascending: false, }, diff --git a/backend/integration_tests/src/client/user_index.rs b/backend/integration_tests/src/client/user_index.rs index 7b5fa4d792..b33cf558de 100644 --- a/backend/integration_tests/src/client/user_index.rs +++ b/backend/integration_tests/src/client/user_index.rs @@ -3,6 +3,7 @@ use user_index_canister::*; // Queries generate_query_call!(check_username); +generate_query_call!(chit_balances); generate_query_call!(current_user); generate_query_call!(search); generate_query_call!(platform_moderators); @@ -31,8 +32,10 @@ generate_update_call!(upgrade_user_canister_wasm); pub mod happy_path { use candid::Principal; use pocket_ic::PocketIc; + use std::collections::HashMap; use types::{ CanisterId, CanisterWasm, Cryptocurrency, DiamondMembershipFees, DiamondMembershipPlanDuration, Empty, UserId, + UserSummary, }; use user_index_canister::users::UserGroup; @@ -103,6 +106,23 @@ pub mod happy_path { } } + pub fn user(env: &PocketIc, canister_id: CanisterId, user_id: UserId) -> UserSummary { + let response = super::user( + env, + Principal::anonymous(), + canister_id, + &user_index_canister::user::Args { + user_id: Some(user_id), + username: None, + }, + ); + + match response { + user_index_canister::user::Response::Success(result) => result, + _ => panic!("User not found"), + } + } + pub fn users( env: &PocketIc, sender: Principal, @@ -212,4 +232,23 @@ pub mod happy_path { user_index_canister::add_platform_operator::Response::Success => {} } } + + pub fn chit_balances( + env: &PocketIc, + user_index_canister_id: CanisterId, + users: Vec, + year: u16, + month: u8, + ) -> HashMap { + let response = super::chit_balances( + env, + Principal::anonymous(), + user_index_canister_id, + &user_index_canister::chit_balances::Args { users, year, month }, + ); + + match response { + user_index_canister::chit_balances::Response::Success(result) => result.balances, + } + } } diff --git a/backend/libraries/types/can.did b/backend/libraries/types/can.did index adc0722f53..c00778f263 100644 --- a/backend/libraries/types/can.did +++ b/backend/libraries/types/can.did @@ -1074,6 +1074,7 @@ type UserSummary = record { suspended : bool; diamond_member : bool; diamond_membership_status : DiamondMembershipStatus; + total_chit_earned : int32; chit_balance : int32; streak : nat16; is_unique_person : bool; @@ -1877,6 +1878,7 @@ type UserSummaryStable = record { }; type UserSummaryVolatile = record { + total_chit_earned : int32; chit_balance : int32; streak : nat16; }; diff --git a/backend/libraries/types/src/user_summary.rs b/backend/libraries/types/src/user_summary.rs index b18ee26e47..bc1c7ebebd 100644 --- a/backend/libraries/types/src/user_summary.rs +++ b/backend/libraries/types/src/user_summary.rs @@ -12,6 +12,7 @@ pub struct UserSummary { pub suspended: bool, pub diamond_member: bool, pub diamond_membership_status: DiamondMembershipStatus, + pub total_chit_earned: i32, pub chit_balance: i32, pub streak: u16, pub is_unique_person: bool, @@ -37,6 +38,7 @@ pub struct UserSummaryStable { #[derive(CandidType, Serialize, Deserialize, Debug)] pub struct UserSummaryVolatile { + pub total_chit_earned: i32, pub chit_balance: i32, pub streak: u16, } diff --git a/backend/libraries/utils/src/time.rs b/backend/libraries/utils/src/time.rs index 53a781ba5b..38e31dbaba 100644 --- a/backend/libraries/utils/src/time.rs +++ b/backend/libraries/utils/src/time.rs @@ -1,3 +1,7 @@ +use candid::Deserialize; +use serde::Serialize; +use std::ops::Range; +use time::{OffsetDateTime, Time}; use types::{Milliseconds, TimestampMillis, TimestampNanos}; pub const SECOND_IN_MS: Milliseconds = 1000; @@ -25,9 +29,92 @@ pub fn tomorrow(now: TimestampMillis) -> TimestampMillis { } pub fn to_date(ts: TimestampMillis) -> time::Date { - time::OffsetDateTime::from_unix_timestamp((ts / 1000) as i64).unwrap().date() + OffsetDateTime::from_unix_timestamp((ts / 1000) as i64).unwrap().date() } pub fn to_timestamp(date: time::Date) -> TimestampMillis { - (time::OffsetDateTime::new_utc(date, time::Time::MIDNIGHT).unix_timestamp() * 1000) as u64 + (OffsetDateTime::new_utc(date, Time::MIDNIGHT).unix_timestamp() * 1000) as u64 +} + +#[derive(Serialize, Deserialize, Copy, Clone, Eq, PartialEq, Ord, PartialOrd)] +pub struct MonthKey { + year: u32, + month: u8, +} + +impl MonthKey { + pub const fn new(year: u32, month: u8) -> MonthKey { + MonthKey { year, month } + } + + pub fn from_timestamp(ts: TimestampMillis) -> MonthKey { + let date = OffsetDateTime::from_unix_timestamp((ts / 1000) as i64).unwrap(); + + MonthKey { + year: date.year() as u32, + month: u8::from(date.month()), + } + } + + pub fn year(&self) -> u32 { + self.year + } + + pub fn month(&self) -> u8 { + self.month + } + + pub fn next(self) -> MonthKey { + if self.month == 12 { + MonthKey { + year: self.year + 1, + month: 1, + } + } else { + MonthKey { + year: self.year, + month: self.month + 1, + } + } + } + + pub fn previous(self) -> MonthKey { + if self.month == 1 { + MonthKey { + year: self.year.saturating_sub(1), + month: 12, + } + } else { + MonthKey { + year: self.year, + month: self.month.saturating_sub(1), + } + } + } + + pub fn timestamp_range(&self) -> Range { + let start = self.start_timestamp(); + let end = self.next().start_timestamp(); + + start..end + } + + pub fn start_timestamp(&self) -> TimestampMillis { + let date = time::Date::from_calendar_date(self.year as i32, self.month.try_into().unwrap(), 1).unwrap(); + (OffsetDateTime::new_utc(date, Time::MIDNIGHT).unix_timestamp() * 1000) as TimestampMillis + } +} + +#[cfg(test)] +mod tests { + use crate::time::MonthKey; + + #[test] + fn month_timestamp_range() { + let month = MonthKey::new(2021, 5); + let range = month.timestamp_range(); + + assert_eq!(range.start, 1619827200000); + assert_eq!(range.end, 1622505600000); + } }