From dee418f47c750c53379169d66f28bedc94e02324 Mon Sep 17 00:00:00 2001 From: X Date: Sun, 15 Feb 2026 10:14:14 +0100 Subject: [PATCH 01/10] blob deletion add blob deletion space reclamation tests formatting test fix minor test fix minor --- e2e/test2.spec.ts | 195 ++++++++++++++++++++++++++++++++---- src/backend/env/mod.rs | 6 ++ src/backend/env/storage.rs | 55 ++++++++-- src/backend/updates.rs | 5 +- src/bucket/src/lib.rs | 107 ++++++++++++++++++-- src/frontend/src/index.html | 4 +- 6 files changed, 333 insertions(+), 39 deletions(-) diff --git a/e2e/test2.spec.ts b/e2e/test2.spec.ts index 691154b0..e5f62277 100644 --- a/e2e/test2.spec.ts +++ b/e2e/test2.spec.ts @@ -87,19 +87,6 @@ test.describe("Regular users flow", () => { await page.getByRole("button", { name: "POST" }).click(); await page.locator("textarea").fill("Hello world!"); - const imagePath = resolve( - __dirname, - "..", - "src", - "frontend", - "assets", - "apple-touch-icon.png", - ); - const [fileChooser] = await Promise.all([ - page.waitForEvent("filechooser"), - page.getByTestId("file-picker").click(), - ]); - await fileChooser.setFiles([imagePath]); await page.getByRole("button", { name: "SUBMIT" }).click(); await page.waitForURL(/#\/post\//); await waitForUILoading(page); @@ -107,9 +94,6 @@ test.describe("Regular users flow", () => { await expect( page.locator("article", { hasText: /Hello world/ }), ).toBeVisible(); - await expect( - page.getByRole("img", { name: "512x512, 2kb" }), - ).toBeVisible(); await page.getByTestId("post-info-toggle").click(); const editButton = page.locator("button[title=Edit]"); @@ -133,15 +117,12 @@ test.describe("Regular users flow", () => { await expect( article.getByText(/Edit:.*this is a post-scriptum/), ).toBeVisible(); - await expect( - page.getByRole("img", { name: "512x512, 2kb" }), - ).toBeVisible(); }); test("Wallet", async () => { await page.getByTestId("toggle-user-section").click(); - await expect(page.getByTestId("credits-balance")).toHaveText("976"); + await expect(page.getByTestId("credits-balance")).toHaveText("996"); await handleDialog( page, @@ -152,7 +133,7 @@ test.describe("Regular users flow", () => { }, ); await waitForUILoading(page); - await expect(page.getByTestId("credits-balance")).toHaveText("2,976"); + await expect(page.getByTestId("credits-balance")).toHaveText("2,996"); const icpBalance = parseFloat( await page.getByTestId("icp-balance").textContent(), @@ -331,4 +312,176 @@ test.describe("Regular users flow", () => { "WONDERLAND", ); }); + + test("Upload image", async () => { + // Navigate to home to clear any overlays from previous test + await page.goto("/"); + await waitForUILoading(page); + + // Sign out bob and sign back in as alice + await page.getByTestId("toggle-user-section").click(); + await page.getByRole("link", { name: /.*SIGN OUT.*/ }).click(); + await waitForUILoading(page); + await page.getByRole("button", { name: "SIGN IN" }).click(); + await page.getByRole("button", { name: "SEED PHRASE" }).click(); + await page + .getByPlaceholder("Enter your seed phrase...") + .fill(mkPwd("alice")); + await page.getByRole("button", { name: "CONTINUE" }).click(); + await waitForUILoading(page); + + // Read credits balance before upload + await page.getByTestId("toggle-user-section").click(); + const balanceBefore = parseInt( + (await page.getByTestId("credits-balance").textContent())!.replace( + /,/g, + "", + ), + ); + await page.getByTestId("toggle-user-section").click(); + + const imagePath = resolve( + __dirname, + "..", + "src", + "frontend", + "assets", + "apple-touch-icon.png", + ); + await page.getByRole("button", { name: "POST" }).click(); + await page.locator("textarea").fill("Post with image"); + const [fileChooser] = await Promise.all([ + page.waitForEvent("filechooser"), + page.getByTestId("file-picker").click(), + ]); + await fileChooser.setFiles([imagePath]); + await page.getByRole("button", { name: "SUBMIT" }).click(); + await page.waitForURL(/#\/post\//); + await waitForUILoading(page); + + await expect( + page.locator("article", { hasText: /Post with image/ }), + ).toBeVisible(); + await expect( + page.getByRole("img", { name: "512x512, 2kb" }), + ).toBeVisible(); + + // Verify credits charged: post_cost (2) + blob_cost (20) = 22 + await page.goto("/"); + await waitForUILoading(page); + await page.getByTestId("toggle-user-section").click(); + const balanceAfter = parseInt( + (await page.getByTestId("credits-balance").textContent())!.replace( + /,/g, + "", + ), + ); + expect(balanceBefore - balanceAfter).toBe(22); + }); + + test("Blob space reclamation", async () => { + const largeImage = resolve( + __dirname, + "..", + "src", + "frontend", + "assets", + "logo.png", + ); + const smallImage = resolve( + __dirname, + "..", + "src", + "frontend", + "assets", + "apple-touch-icon.png", + ); + + const getImageOffset = async (altText: string): Promise => { + const img = page.getByRole("img", { name: altText }); + await expect(img).toBeVisible(); + const src = await img.getAttribute("src"); + const match = src!.match(/offset=(\d+)/); + return parseInt(match![1]); + }; + + const uploadPost = async ( + text: string, + imagePath: string, + ): Promise => { + await page.getByRole("button", { name: "POST" }).click(); + await page.locator("textarea").fill(text); + const [fileChooser] = await Promise.all([ + page.waitForEvent("filechooser"), + page.getByTestId("file-picker").click(), + ]); + await fileChooser.setFiles([imagePath]); + await page.getByRole("button", { name: "SUBMIT" }).click(); + await page.waitForURL(/#\/post\//); + await waitForUILoading(page); + }; + + // Read initial credits balance + await page.goto("/"); + await waitForUILoading(page); + await page.getByTestId("toggle-user-section").click(); + const initialBalance = parseInt( + (await page.getByTestId("credits-balance").textContent())!.replace( + /,/g, + "", + ), + ); + await page.getByTestId("toggle-user-section").click(); + + // Create post 1 with large image + await uploadPost("Large image post", largeImage); + const offset1 = await getImageOffset("200x200, 6kb"); + const post1Url = page.url(); + + // Create post 2 with small image + await page.goto("/"); + await waitForUILoading(page); + await uploadPost("Small image post", smallImage); + const offset2 = await getImageOffset("512x512, 2kb"); + + // Post 2 should be after post 1 in storage + expect(offset2).toBeGreaterThan(offset1); + + // Delete post 1 + await page.goto(post1Url); + await waitForUILoading(page); + await page + .locator(".post_box", { hasText: /Large image post/ }) + .getByTestId("post-info-toggle") + .click(); + await handleDialog(page, "confirm the post deletion", "", async () => { + await page.locator("button[title='Delete post']").click(); + }); + await waitForUILoading(page); + + // Wait for the async free call to complete on the bucket + await page.waitForTimeout(3000); + + // Create post 3 with large image (same size as post 1) + await page.goto("/"); + await waitForUILoading(page); + await uploadPost("Reclaimed space post", largeImage); + const offset3 = await getImageOffset("200x200, 6kb"); + + // Post 3 should have reclaimed the space freed by post 1 + expect(offset3).toBe(offset1); + expect(offset3).toBeLessThan(offset2); + + // Verify total credits charged: 3 uploads (3*22=66) + 1 deletion (2) = 68 + await page.goto("/"); + await waitForUILoading(page); + await page.getByTestId("toggle-user-section").click(); + const finalBalance = parseInt( + (await page.getByTestId("credits-balance").textContent())!.replace( + /,/g, + "", + ), + ); + expect(initialBalance - finalBalance).toBe(68); + }); }); diff --git a/src/backend/env/mod.rs b/src/backend/env/mod.rs index 9a8bcecd..6e48f371 100644 --- a/src/backend/env/mod.rs +++ b/src/backend/env/mod.rs @@ -5,6 +5,7 @@ use self::invoices::{ICPInvoice, USER_ICP_SUBACCOUNT}; use self::post::{archive_cold_posts, Extension, Post, PostId}; use self::post_iterators::{IteratorMerger, MergeStrategy}; use self::proposals::{Payload, ReleaseInfo, Status}; +use self::storage::Storage; use self::token::{account, TransferArgs}; use self::user::{Filters, Mode, Notification, Predicate, UserFilter}; use crate::assets::export_token_supply; @@ -2722,12 +2723,17 @@ impl State { _ => {} }; + let files = post.files.clone(); Post::mutate(self, &post_id, |post| { post.delete(versions.clone()); Ok(()) }) .expect("couldn't delete post"); + if !files.is_empty() { + ic_cdk::spawn(Storage::free_blobs(files)); + } + Ok(()) } diff --git a/src/backend/env/storage.rs b/src/backend/env/storage.rs index c9cc5895..3d79b306 100644 --- a/src/backend/env/storage.rs +++ b/src/backend/env/storage.rs @@ -18,19 +18,20 @@ const BUCKET_WASM_GZ: &[u8] = include_bytes!("../../../target/wasm32-unknown-unknown/release/bucket.wasm.gz"); impl Storage { - async fn allocate_space() -> Result { - if let Some(id) = read(|state| { + async fn allocate_space() -> Result<(Principal, u64), String> { + if let Some((id, offset)) = read(|state| { state .storage .buckets .iter() - .find_map(|(id, size)| (*size < CONFIG.max_bucket_size).then_some(*id)) + .find_map(|(id, size)| (*size < CONFIG.max_bucket_size).then_some((*id, *size))) }) { - return Ok(id); + return Ok((id, offset)); } let id = crate::canisters::new().await?; + let init_offset = 8; mutate(|state| { - state.storage.buckets.insert(id, 0); + state.storage.buckets.insert(id, init_offset); state.logger.debug(format!("New bucket {} created.", id)); }); canisters::install(id, BUCKET_WASM_GZ, CanisterInstallMode::Install).await?; @@ -39,20 +40,58 @@ impl Storage { .logger .debug(format!("WASM installed to bucket {}.", id)); }); - Ok(id) + Ok((id, init_offset)) } pub async fn write_to_bucket(blob: &[u8]) -> Result<(Principal, u64), String> { - let id = Storage::allocate_space().await?; + let (id, curr_offset) = Storage::allocate_space().await?; let response = canisters::call_canister_raw(id, "write", blob) .await .map_err(|err| format!("couldn't call write on a bucket: {:?}", err))?; let mut offset_bytes: [u8; 8] = Default::default(); offset_bytes.copy_from_slice(&response); let offset = u64::from_be_bytes(offset_bytes); - mutate(|state| state.storage.buckets.insert(id, offset + blob.len() as u64)); + let new_offset = offset + blob.len() as u64; + mutate(|state| { + state + .storage + .buckets + .insert(id, curr_offset.max(new_offset)) + }); Ok((id, offset)) } + + /// Frees blobs on their respective bucket canisters. + /// The `files` map has keys in the format `"blob_id@bucket_principal"` + /// and values of `(offset, length)`. + pub async fn free_blobs(files: BTreeMap) { + // Group blobs by bucket canister. + let mut by_bucket: BTreeMap> = BTreeMap::new(); + for (key, (offset, length)) in &files { + if let Some(bucket_id) = key + .rsplit_once('@') + .and_then(|(_, b)| b.parse::().ok()) + { + by_bucket + .entry(bucket_id) + .or_default() + .push((*offset, *length as u64)); + } + } + + for (bucket_id, segments) in by_bucket { + if let Err(err) = + canisters::call_canister::<_, ()>(bucket_id, "free", (segments,)).await + { + mutate(|state| { + state.logger.error(format!( + "couldn't free blobs on bucket {}: {:?}", + bucket_id, err + )) + }); + } + } + } } #[allow(dead_code)] diff --git a/src/backend/updates.rs b/src/backend/updates.rs index 0eb157be..5d5a2677 100644 --- a/src/backend/updates.rs +++ b/src/backend/updates.rs @@ -11,6 +11,7 @@ use env::{ canisters::get_full_neuron, config::CONFIG, post::{Extension, Post, PostId}, + storage, user::{Draft, User, UserId}, State, }; @@ -112,7 +113,9 @@ fn post_upgrade() { fn sync_post_upgrade_fixtures() {} #[allow(clippy::all)] -async fn async_post_upgrade_fixtures() {} +async fn async_post_upgrade_fixtures() { + storage::upgrade_buckets().await; +} /* * UPDATES diff --git a/src/bucket/src/lib.rs b/src/bucket/src/lib.rs index 6f39008e..07b404b9 100644 --- a/src/bucket/src/lib.rs +++ b/src/bucket/src/lib.rs @@ -1,4 +1,4 @@ -use candid::{CandidType, Deserialize, Principal}; +use candid::{CandidType, Decode, Deserialize, Encode, Principal}; use ic_cdk::api::{ self, call::{arg_data_raw, reply_raw}, @@ -7,6 +7,7 @@ use ic_cdk::api::{ }; use serde::Serialize; use serde_bytes::ByteBuf; +use std::cell::RefCell; mod url; @@ -14,9 +15,22 @@ mod url; // without trying to read the stable memory. const MAX_BLOB_SIZE: u64 = 8 * 1024 * 1024; +// Minimum remainder size (50KB) to keep when splitting a free segment. +const MIN_REMAINDER: u64 = 50 * 1024; + // HTTP request and response headers. type Headers = Vec<(String, String)>; +#[derive(Clone, Copy, CandidType, Deserialize)] +struct Segment { + start: u64, + length: u64, +} + +thread_local! { + static FREE_LIST: RefCell> = const { RefCell::new(Vec::new()) }; +} + #[derive(CandidType, Deserialize)] struct HttpRequest { url: String, @@ -51,9 +65,51 @@ fn init() { set_controller(); } +#[export_name = "canister_pre_upgrade"] +fn pre_upgrade() { + let offset = read_offset(); + + FREE_LIST.with(|fl| { + let bytes = Encode!(&*fl.borrow()).expect("couldn't serialize free list"); + let len = bytes.len() as u64; + grow_to_fit(offset, 8 + len); + stable_write(offset, &len.to_be_bytes()); + stable_write(offset + 8, &bytes); + }); +} + #[export_name = "canister_post_upgrade"] fn post_upgrade() { set_controller(); + + let offset = read_offset(); + let stable_mem_size = stable_size() << 16; + // Not enough stable memory to even read the free-list length header. + if offset + 8 > stable_mem_size { + return; + } + + let mut len_bytes: [u8; 8] = Default::default(); + stable_read(offset, &mut len_bytes); + let len = u64::from_be_bytes(len_bytes); + // On first upgrade from old code (no pre_upgrade), stable memory past the + // high-water mark is zero-initialized, so len will be 0 and we skip gracefully. + if len == 0 || offset + 8 + len > stable_mem_size { + return; + } + + let mut bytes = vec![0u8; len as usize]; + stable_read(offset + 8, &mut bytes); + + if let Ok(free_list) = Decode!(&bytes, Vec) { + FREE_LIST.with(|fl| *fl.borrow_mut() = free_list); + } +} + +fn read_offset() -> u64 { + let mut bytes: [u8; 8] = Default::default(); + api::stable::stable_read(0, &mut bytes); + u64::from_be_bytes(bytes) } #[ic_cdk_macros::query] @@ -114,17 +170,54 @@ fn http_image(args: &str) -> HttpResponse { #[export_name = "canister_update write"] fn write() { assert_controller(); - let mut offset_bytes: [u8; 8] = Default::default(); - api::stable::stable_read(0, &mut offset_bytes); let blob = arg_data_raw(); - let offset = u64::from_be_bytes(offset_bytes); - grow_to_fit(offset, blob.len() as u64); + let blob_len = blob.len() as u64; + + let offset = FREE_LIST.with(|fl| { + let mut free_list = fl.borrow_mut(); + + // Binary-search for the smallest free segment that fits. + // Ok(i) = exact match, Err(i) = next larger segment at index i. + match free_list.binary_search_by_key(&blob_len, |s| s.length) { + Ok(idx) | Err(idx) if idx < free_list.len() => { + let seg = free_list.remove(idx); + let remainder = seg.length - blob_len; + if remainder >= MIN_REMAINDER { + free_list.push(Segment { + start: seg.start + blob_len, + length: remainder, + }); + free_list.sort_by_key(|s| s.length); + } + seg.start + } + _ => { + // All free segments are too small; append at the end. + let offset = read_offset(); + grow_to_fit(offset, blob_len); + let new_offset = offset + blob_len; + api::stable::stable_write(0, &new_offset.to_be_bytes()); + offset + } + } + }); + stable_write(offset, &blob); - let new_offset = offset + blob.len() as u64; - api::stable::stable_write(0, &new_offset.to_be_bytes()); reply_raw(&offset.to_be_bytes()); } +#[ic_cdk_macros::update] +fn free(segments: Vec<(u64, u64)>) { + assert_controller(); + FREE_LIST.with(|fl| { + let mut free_list = fl.borrow_mut(); + for (start, length) in segments { + free_list.push(Segment { start, length }); + } + free_list.sort_by_key(|s| s.length); + }); +} + fn grow_to_fit(offset: u64, len: u64) { if offset + len < (stable_size() << 16) { return; diff --git a/src/frontend/src/index.html b/src/frontend/src/index.html index 58bd50c2..9f217766 100644 --- a/src/frontend/src/index.html +++ b/src/frontend/src/index.html @@ -58,7 +58,7 @@ src: url("font-regular.woff2") format("woff2"); font-weight: normal; font-style: normal; - font-display: swap; + font-display: optional; } @font-face { @@ -66,7 +66,7 @@ src: url("font-bold.woff2") format("woff2"); font-weight: bold; font-style: normal; - font-display: swap; + font-display: optional; }