Merge pull request #13 from NeutronStarDAO/bug-feed-canister-batch_re…

…ceive_feed-func Fix feed canister batch_receive_feed func bug
NeutronStarDAO · Aug 15, 2024 · e9cd156 · e9cd156
2 parents f6509c1 + c2599dc
commit e9cd156
Show file tree

Hide file tree

Showing 14 changed files with 241 additions and 7 deletions.
diff --git a/feed/feed/feed.did b/feed/feed/feed.did
@@ -54,6 +54,10 @@ service : (principal, principal, principal) -> {
   get_feed_number : (principal) -> (nat64) query;
   get_latest_feed : (principal, nat64) -> (vec Post) query;
   get_post : (text) -> (opt Post) query;
+  get_post_fetch : () -> (principal) query;
+  get_post_index : () -> (nat64) query;
   get_post_number : (principal) -> (nat64) query;
+  get_root_bucket : () -> (principal) query;
+  get_user_actor : () -> (principal) query;
   status : () -> (CanisterStatusResponse);
 }
diff --git a/feed/feed/src/lib.rs b/feed/feed/src/lib.rs
@@ -427,7 +427,11 @@ async fn batch_receive_feed(
             bucket, 
             "get_post", 
             (post_id.clone(), )
-        ).await.unwrap().0.unwrap();
+        ).await.unwrap().0;
+
+        if let None = new_post {
+            continue;
+        }
 
         let user_map = FEED_MAP.with(|map| {
             map.borrow().get(&user)
@@ -437,13 +441,13 @@ async fn batch_receive_feed(
             None => {
                 // 被通知者第一次收到 feed 推流
                 let mut user_map = FeedHashMap(HashMap::new());
-                user_map.0.insert(post_id.clone(), new_post);
+                user_map.0.insert(post_id.clone(), new_post.unwrap());
                 FEED_MAP.with(|map| {
                     map.borrow_mut().insert(user, user_map);
                 })
             },
             Some(mut user_map) => {
-                user_map.0.insert(post_id.clone(), new_post);
+                user_map.0.insert(post_id.clone(), new_post.unwrap());
                 FEED_MAP.with(|map| {
                     map.borrow_mut().insert(user, user_map);
                 });
@@ -482,6 +486,11 @@ async fn check_available_bucket() -> bool {
     true
 }
 
+#[ic_cdk::query]
+fn get_post_index() -> u64 {
+    POST_INDEX.with(|post_index| post_index.borrow().get().clone())
+}
+
 #[ic_cdk::query]
 fn get_bucket() -> Option<Principal> {
     BUCKET.with(|pr| {
@@ -623,6 +632,21 @@ fn get_all_latest_feed(
     result
 }
 
+#[ic_cdk::query]
+fn get_root_bucket() -> Principal {
+    ROOT_BUCKET.with(|root_bucket| root_bucket.borrow().get().clone())
+}
+
+#[ic_cdk::query]
+fn get_post_fetch() -> Principal {
+    POST_FETCH_ACTOR.with(|post_fetch| post_fetch.borrow().get().clone())
+}
+
+#[ic_cdk::query]
+fn get_user_actor() -> Principal {
+    USER_ACTOR.with(|user| user.borrow().get().clone())
+}
+
 fn get_post_id(canister: &Principal, user: &Principal, index: u64) -> String {
     canister.to_text() + "#" + &user.to_text() + "#" + &index.to_string()   
 }

diff --git a/feed/root_feed/root_feed.did b/feed/root_feed/root_feed.did
@@ -27,10 +27,24 @@ type QueryStats = record {
 };
 service : (principal, principal) -> {
   create_feed_canister : () -> (principal);
+  get_all_feed_canister : () -> (vec principal) query;
+  get_available_feed_canister_index : () -> (nat64) query;
+  get_feed_canister_index : () -> (nat64) query;
+  get_feed_canister_users_number_entries : () -> (
+      vec record { principal; nat64 },
+    ) query;
   get_feed_wasm : () -> (blob) query;
+  get_root_bucket : () -> (principal) query;
+  get_user_actor : () -> (principal) query;
   get_user_feed_canister : (principal) -> (opt principal) query;
+  get_user_feed_canister_entries : () -> (
+      vec record { principal; principal },
+    ) query;
   init_fetch_actor : (principal) -> ();
   init_user_feed : () -> (principal);
+  set_root_bucket : (principal) -> (bool);
+  set_user_actor : (principal) -> (bool);
   status : () -> (CanisterStatusResponse);
+  update_feed_canister_controller : (principal) -> (bool);
   update_feed_wasm : (blob, nat64) -> (bool);
 }
diff --git a/feed/root_feed/src/lib.rs b/feed/root_feed/src/lib.rs
@@ -52,21 +52,21 @@ thread_local! {
 
     static ROOT_BUCKET: RefCell<StableCell<Principal, Memory>> = RefCell::new(
         StableCell::init(
-            MEMORY_MANAGER.with(|m| m.borrow().get(MemoryId::new(3))), 
+            MEMORY_MANAGER.with(|m| m.borrow().get(MemoryId::new(5))), 
             Principal::anonymous()
         ).unwrap() 
     );
 
     static USER_ACTOR: RefCell<StableCell<Principal, Memory>> = RefCell::new(
         StableCell::init(
-            MEMORY_MANAGER.with(|m| m.borrow().get(MemoryId::new(4))), 
+            MEMORY_MANAGER.with(|m| m.borrow().get(MemoryId::new(6))), 
             Principal::anonymous()
         ).unwrap() 
     );
 
     static POST_FETCH_ACTOR: RefCell<StableCell<Principal, Memory>> = RefCell::new(
         StableCell::init(
-            MEMORY_MANAGER.with(|m| m.borrow().get(MemoryId::new(5))), 
+            MEMORY_MANAGER.with(|m| m.borrow().get(MemoryId::new(7))), 
             Principal::anonymous()
         ).unwrap() 
     );
@@ -215,12 +215,135 @@ fn get_user_feed_canister(user: Principal) -> Option<Principal> {
     })
 }
 
+#[ic_cdk::query]
+fn get_feed_canister_index() -> u64 {
+    FEED_CANISTER_INDEX.with(|index| index.borrow().get().clone())
+}
+
+#[ic_cdk::query]
+fn get_available_feed_canister_index() -> u64 {
+    AVAILABLE_FEED_CANISTER_INDEX.with(|index| index.borrow().get().clone())
+}
+
+#[ic_cdk::query] 
+fn get_all_feed_canister() -> Vec<Principal> {
+    FEED_CANISTER_MAP.with(|map| {
+        let mut feed_canister_list = Vec::new();
+        for (_, feed_canister) in map.borrow().iter() {
+            feed_canister_list.push(feed_canister)
+        }
+        feed_canister_list
+    })
+}
+
+#[ic_cdk::query]
+fn get_user_feed_canister_entries() -> Vec<(Principal, Principal)> {
+    USER_FEED_CANISTER.with(|map| {
+        let mut entries = Vec::new();
+
+        for (user, feed) in map.borrow().iter() {
+            entries.push((user, feed))
+        }
+
+        entries
+    })
+}
+
+#[ic_cdk::query]
+fn get_feed_canister_users_number_entries() -> Vec<(Principal, u64)> {
+    FEED_CANISTER_USERS_NUMBER.with(|map| {
+        let mut entries = Vec::new();
+
+        for (feed, number) in map.borrow().iter() {
+            entries.push((feed, number))
+        }
+
+        entries
+    })
+}
+
+#[ic_cdk::query]
+fn get_root_bucket() -> Principal {
+    ROOT_BUCKET.with(|root_bucket| root_bucket.borrow().get().clone())
+}
+
+#[ic_cdk::query]
+fn get_user_actor() -> Principal {
+    USER_ACTOR.with(|user_actor| user_actor.borrow().get().clone())
+}
+
+#[ic_cdk::update]
+async fn set_root_bucket(canister: Principal) -> bool {
+    if !is_controller(&ic_cdk::caller()).await {
+        return false;
+    }
+
+    ROOT_BUCKET.with(|root_bucket| root_bucket.borrow_mut().set(canister).unwrap());
+
+    true
+}
+
+#[ic_cdk::update]
+async fn set_user_actor(canister: Principal) -> bool {
+    if !is_controller(&ic_cdk::caller()).await {
+        return false;
+    }
+
+    USER_ACTOR.with(|user_actor| user_actor.borrow_mut().set(canister).unwrap());
+
+    true
+}
+
 #[ic_cdk::update]
 async fn status() -> CanisterStatusResponse {
     ic_cdk::api::management_canister::main::canister_status(CanisterIdRecord {
         canister_id: ic_cdk::api::id()
     }).await.unwrap().0
 }
 
+
+async fn is_controller(user: &Principal) -> bool {
+    let status = status().await;
+    let controllers = status.settings.controllers;
+
+    if !controllers.contains(user) {
+        return false;
+    }
+
+    true
+}
+
+#[ic_cdk::update]
+async fn update_feed_canister_controller(
+    controller: Principal
+) -> bool {
+    if !is_controller(&ic_cdk::caller()).await {
+        return false;
+    }
+
+    let controllers = vec![ic_cdk::api::id(), controller];
+
+    let feed_canister_vec = get_all_feed_canister();
+
+    for feed_canister in feed_canister_vec {
+        ic_cdk::api::management_canister::main::update_settings(
+            ic_cdk::api::management_canister::main::UpdateSettingsArgument {
+                canister_id: feed_canister,
+                settings: CanisterSettings {
+                    controllers: Some(controllers.clone()),
+                    compute_allocation: None,
+                    memory_allocation: None,
+                    freezing_threshold: None,
+                    reserved_cycles_limit: None,
+                    wasm_memory_limit: None,
+                    log_visibility: None
+                }
+            }
+        ).await.unwrap();
+    }
+
+    true
+}
+
 // Enable Candid export
 ic_cdk::export_candid!();
diff --git a/regenerate_candid.sh b/regenerate_candid.sh
@@ -14,4 +14,6 @@ candid-extractor target/wasm32-unknown-unknown/release/root_fetch.wasm > fetch/r
 
 candid-extractor target/wasm32-unknown-unknown/release/user.wasm > user/user.did
 
-candid-extractor target/wasm32-unknown-unknown/release/photo_storage.wasm > storage/photo_storage/photo_storage.did
+candid-extractor target/wasm32-unknown-unknown/release/photo_storage.wasm > storage/photo_storage/photo_storage.did
+
+dfx generate
diff --git a/src/declarations/feed/feed.did b/src/declarations/feed/feed.did
@@ -54,6 +54,10 @@ service : (principal, principal, principal) -> {
   get_feed_number : (principal) -> (nat64) query;
   get_latest_feed : (principal, nat64) -> (vec Post) query;
   get_post : (text) -> (opt Post) query;
+  get_post_fetch : () -> (principal) query;
+  get_post_index : () -> (nat64) query;
   get_post_number : (principal) -> (nat64) query;
+  get_root_bucket : () -> (principal) query;
+  get_user_actor : () -> (principal) query;
   status : () -> (CanisterStatusResponse);
 }
diff --git a/src/declarations/feed/feed.did.d.ts b/src/declarations/feed/feed.did.d.ts
@@ -65,7 +65,11 @@ export interface _SERVICE {
   'get_feed_number' : ActorMethod<[Principal], bigint>,
   'get_latest_feed' : ActorMethod<[Principal, bigint], Array<Post>>,
   'get_post' : ActorMethod<[string], [] | [Post]>,
+  'get_post_fetch' : ActorMethod<[], Principal>,
+  'get_post_index' : ActorMethod<[], bigint>,
   'get_post_number' : ActorMethod<[Principal], bigint>,
+  'get_root_bucket' : ActorMethod<[], Principal>,
+  'get_user_actor' : ActorMethod<[], Principal>,
   'status' : ActorMethod<[], CanisterStatusResponse>,
 }
 export declare const idlFactory: IDL.InterfaceFactory;

diff --git a/src/declarations/feed/feed.did.js b/src/declarations/feed/feed.did.js
@@ -70,7 +70,11 @@ export const idlFactory = ({ IDL }) => {
         ['query'],
       ),
     'get_post' : IDL.Func([IDL.Text], [IDL.Opt(Post)], ['query']),
+    'get_post_fetch' : IDL.Func([], [IDL.Principal], ['query']),
+    'get_post_index' : IDL.Func([], [IDL.Nat64], ['query']),
     'get_post_number' : IDL.Func([IDL.Principal], [IDL.Nat64], ['query']),
+    'get_root_bucket' : IDL.Func([], [IDL.Principal], ['query']),
+    'get_user_actor' : IDL.Func([], [IDL.Principal], ['query']),
     'status' : IDL.Func([], [CanisterStatusResponse], []),
   });
 };

diff --git a/src/declarations/root_feed/root_feed.did b/src/declarations/root_feed/root_feed.did
@@ -27,10 +27,24 @@ type QueryStats = record {
 };
 service : (principal, principal) -> {
   create_feed_canister : () -> (principal);
+  get_all_feed_canister : () -> (vec principal) query;
+  get_available_feed_canister_index : () -> (nat64) query;
+  get_feed_canister_index : () -> (nat64) query;
+  get_feed_canister_users_number_entries : () -> (
+      vec record { principal; nat64 },
+    ) query;
   get_feed_wasm : () -> (blob) query;
+  get_root_bucket : () -> (principal) query;
+  get_user_actor : () -> (principal) query;
   get_user_feed_canister : (principal) -> (opt principal) query;
+  get_user_feed_canister_entries : () -> (
+      vec record { principal; principal },
+    ) query;
   init_fetch_actor : (principal) -> ();
   init_user_feed : () -> (principal);
+  set_root_bucket : (principal) -> (bool);
+  set_user_actor : (principal) -> (bool);
   status : () -> (CanisterStatusResponse);
+  update_feed_canister_controller : (principal) -> (bool);
   update_feed_wasm : (blob, nat64) -> (bool);
 }
diff --git a/src/declarations/root_feed/root_feed.did.d.ts b/src/declarations/root_feed/root_feed.did.d.ts
@@ -34,11 +34,27 @@ export interface QueryStats {
 }
 export interface _SERVICE {
   'create_feed_canister' : ActorMethod<[], Principal>,
+  'get_all_feed_canister' : ActorMethod<[], Array<Principal>>,
+  'get_available_feed_canister_index' : ActorMethod<[], bigint>,
+  'get_feed_canister_index' : ActorMethod<[], bigint>,
+  'get_feed_canister_users_number_entries' : ActorMethod<
+    [],
+    Array<[Principal, bigint]>
+  >,
   'get_feed_wasm' : ActorMethod<[], Uint8Array | number[]>,
+  'get_root_bucket' : ActorMethod<[], Principal>,
+  'get_user_actor' : ActorMethod<[], Principal>,
   'get_user_feed_canister' : ActorMethod<[Principal], [] | [Principal]>,
+  'get_user_feed_canister_entries' : ActorMethod<
+    [],
+    Array<[Principal, Principal]>
+  >,
   'init_fetch_actor' : ActorMethod<[Principal], undefined>,
   'init_user_feed' : ActorMethod<[], Principal>,
+  'set_root_bucket' : ActorMethod<[Principal], boolean>,
+  'set_user_actor' : ActorMethod<[Principal], boolean>,
   'status' : ActorMethod<[], CanisterStatusResponse>,
+  'update_feed_canister_controller' : ActorMethod<[Principal], boolean>,
   'update_feed_wasm' : ActorMethod<[Uint8Array | number[], bigint], boolean>,
 }
 export declare const idlFactory: IDL.InterfaceFactory;

diff --git a/src/declarations/root_feed/root_feed.did.js b/src/declarations/root_feed/root_feed.did.js
@@ -35,15 +35,37 @@ export const idlFactory = ({ IDL }) => {
   });
   return IDL.Service({
     'create_feed_canister' : IDL.Func([], [IDL.Principal], []),
+    'get_all_feed_canister' : IDL.Func([], [IDL.Vec(IDL.Principal)], ['query']),
+    'get_available_feed_canister_index' : IDL.Func([], [IDL.Nat64], ['query']),
+    'get_feed_canister_index' : IDL.Func([], [IDL.Nat64], ['query']),
+    'get_feed_canister_users_number_entries' : IDL.Func(
+        [],
+        [IDL.Vec(IDL.Tuple(IDL.Principal, IDL.Nat64))],
+        ['query'],
+      ),
     'get_feed_wasm' : IDL.Func([], [IDL.Vec(IDL.Nat8)], ['query']),
+    'get_root_bucket' : IDL.Func([], [IDL.Principal], ['query']),
+    'get_user_actor' : IDL.Func([], [IDL.Principal], ['query']),
     'get_user_feed_canister' : IDL.Func(
         [IDL.Principal],
         [IDL.Opt(IDL.Principal)],
         ['query'],
       ),
+    'get_user_feed_canister_entries' : IDL.Func(
+        [],
+        [IDL.Vec(IDL.Tuple(IDL.Principal, IDL.Principal))],
+        ['query'],
+      ),
     'init_fetch_actor' : IDL.Func([IDL.Principal], [], []),
     'init_user_feed' : IDL.Func([], [IDL.Principal], []),
+    'set_root_bucket' : IDL.Func([IDL.Principal], [IDL.Bool], []),
+    'set_user_actor' : IDL.Func([IDL.Principal], [IDL.Bool], []),
     'status' : IDL.Func([], [CanisterStatusResponse], []),
+    'update_feed_canister_controller' : IDL.Func(
+        [IDL.Principal],
+        [IDL.Bool],
+        [],
+      ),
     'update_feed_wasm' : IDL.Func(
         [IDL.Vec(IDL.Nat8), IDL.Nat64],
         [IDL.Bool],