Rust library `canister_tools` for simple upgrades, data-snapshots download and upload, and some handy tools

Developers Rust
3

Hi @ydr990318, sure thing. From your post, it looks like your question is about changing the type or adding or removing fields of the global data during a canister upgrade.

The first step to integrate the library is to add it to the Cargo.toml file of your project, in the dependencies section, canister-tools = "0.2.2". Then, global variables in the main memory of a rust canister are stored in a thread_local with a refcell like this:

#[derive(Default, CandidType, Deserialize)]
struct Data {
    field_one: u64,
    field_two: String,
}

thread_local! {
    static DATA: RefCell<Data> = RefCell::new(Data::default());
}

In your code, you can access this data using the with_borrow and with_borrow_mut methods specified here: LocalKey in std::thread - Rust like this:

#[update]
fn sample_canister_method() {
    DATA.with_borrow_mut(|data| {
        data.field_one += 1;
    });
}

To save this data through upgrades, first choose a stable-memory memory-id for this data-type. There can be memory-ids from 0-255 and make sure to use a different memory-id for each global-variable you have. Here we have one global variable and we choose memory-id: 0.

use canister_tools::MemoryId;
const DATA_MEMORY_ID: MemoryId = MemoryId::new(0);

Call the canister_tools::init, canister_tools::pre_upgrade, and canister_tools::post_upgrade functions in the canister init and upgrade hooks like this:

#[init]
fn init() {
    canister_tools::init(&DATA, DATA_MEMORY_ID);
}  
  
#[pre_upgrade]
fn pre_upgrade() {
    canister_tools::pre_upgrade();
}

#[post_upgrade]
fn post_upgrade() {
    canister_tools::post_upgrade(&DATA, DATA_MEMORY_ID, None::<fn(Data) -> Data>);
}

Now the data will persist through upgrades.

When it comes time to change the Data type, adding or removing some fields, you can use the built-in opt_old_as_new_convert parameter on the canister_tools::post_upgrade function. First we keep the existing Data type definition in the canister for the upgrade and rename it to OldData. Then define the new Data type next to it like this:

#[derive(Default, CandidType, Deserialize)]
struct OldData {
    field_one: u64,
    field_two: String,
}

#[derive(Default, CandidType, Deserialize)]
struct Data {
    field_one: u64,
    field_two: String,
    new_field: u64,
}

Then in the post_upgrade function, pass a conversion function from the OldData to the new Data like this:

#[post_upgrade]
fn post_upgrade() {
    canister_tools::post_upgrade(
        &DATA, 
        DATA_MEMORY_ID, 
        Some::<fn(OldData) -> Data>(
            |old_data: OldData| {
                Data{
                    field_one: old_data.field_one,
                    field_two: old_data.field_two,
                    new_field: 564, // set the new_field value
                }
            }
        )
    );
}

Then upgrade the canister, and after the upgrade, remove the OldData type definition from your code, and set the opt_old_as_new_convert parameter back to None in the canister_tools::post_upgrade function:

#[post_upgrade]
fn post_upgrade() {
    canister_tools::post_upgrade(&DATA, DATA_MEMORY_ID, None::<fn(Data) -> Data>);
}

And there you go.