Skip to content

Commit 65357d3

Browse files
LandynDevLandynDev
committed
Address review: remove redundant clamping, add atomicity
1 parent 1c4f180 commit 65357d3

2 files changed

Lines changed: 197 additions & 70 deletions

File tree

‎chain-extensions/src/lib.rs‎

Lines changed: 55 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ pub mod types;
99

1010
use crate::types::{FunctionId, Output};
1111
use codec::{Decode, Encode, MaxEncodedLen};
12+
use frame_support::storage::{TransactionOutcome, transactional};
1213
use frame_support::{DebugNoBound, traits::Get};
1314
use frame_system::RawOrigin;
1415
use pallet_contracts::chain_extension::{
@@ -535,22 +536,16 @@ where
535536

536537
let caller = env.caller();
537538

538-
let alpha_available =
539-
pallet_subtensor::Pallet::<T>::get_stake_for_hotkey_and_coldkey_on_subnet(
540-
&hotkey, &caller, netuid,
541-
);
542-
let actual_amount = amount.min(alpha_available);
543-
544539
let call_result = pallet_subtensor::Pallet::<T>::recycle_alpha(
545540
RawOrigin::Signed(caller).into(),
546541
hotkey,
547-
actual_amount,
542+
amount,
548543
netuid,
549544
);
550545

551546
match call_result {
552547
Ok(_) => {
553-
env.write_output(&actual_amount.encode())
548+
env.write_output(&amount.encode())
554549
.map_err(|_| DispatchError::Other("Failed to write output"))?;
555550
Ok(RetVal::Converging(Output::Success as u32))
556551
}
@@ -572,22 +567,16 @@ where
572567

573568
let caller = env.caller();
574569

575-
let alpha_available =
576-
pallet_subtensor::Pallet::<T>::get_stake_for_hotkey_and_coldkey_on_subnet(
577-
&hotkey, &caller, netuid,
578-
);
579-
let actual_amount = amount.min(alpha_available);
580-
581570
let call_result = pallet_subtensor::Pallet::<T>::burn_alpha(
582571
RawOrigin::Signed(caller).into(),
583572
hotkey,
584-
actual_amount,
573+
amount,
585574
netuid,
586575
);
587576

588577
match call_result {
589578
Ok(_) => {
590-
env.write_output(&actual_amount.encode())
579+
env.write_output(&amount.encode())
591580
.map_err(|_| DispatchError::Other("Failed to write output"))?;
592581
Ok(RetVal::Converging(Output::Success as u32))
593582
}
@@ -609,33 +598,33 @@ where
609598

610599
let caller = env.caller();
611600

612-
let alpha = pallet_subtensor::Pallet::<T>::do_add_stake(
613-
RawOrigin::Signed(caller.clone()).into(),
614-
hotkey.clone(),
615-
netuid,
616-
tao_amount,
617-
);
601+
let result = transactional::with_transaction(|| {
602+
let alpha = match pallet_subtensor::Pallet::<T>::do_add_stake(
603+
RawOrigin::Signed(caller.clone()).into(),
604+
hotkey.clone(),
605+
netuid,
606+
tao_amount,
607+
) {
608+
Ok(a) => a,
609+
Err(e) => return TransactionOutcome::Rollback(Err(e)),
610+
};
611+
612+
match pallet_subtensor::Pallet::<T>::recycle_alpha(
613+
RawOrigin::Signed(caller).into(),
614+
hotkey,
615+
alpha,
616+
netuid,
617+
) {
618+
Ok(_) => TransactionOutcome::Commit(Ok(alpha)),
619+
Err(e) => TransactionOutcome::Rollback(Err(e)),
620+
}
621+
});
618622

619-
match alpha {
623+
match result {
620624
Ok(alpha) => {
621-
let recycle_result = pallet_subtensor::Pallet::<T>::recycle_alpha(
622-
RawOrigin::Signed(caller).into(),
623-
hotkey,
624-
alpha,
625-
netuid,
626-
);
627-
628-
match recycle_result {
629-
Ok(_) => {
630-
env.write_output(&alpha.encode())
631-
.map_err(|_| DispatchError::Other("Failed to write output"))?;
632-
Ok(RetVal::Converging(Output::Success as u32))
633-
}
634-
Err(e) => {
635-
let error_code = Output::from(e) as u32;
636-
Ok(RetVal::Converging(error_code))
637-
}
638-
}
625+
env.write_output(&alpha.encode())
626+
.map_err(|_| DispatchError::Other("Failed to write output"))?;
627+
Ok(RetVal::Converging(Output::Success as u32))
639628
}
640629
Err(e) => {
641630
let error_code = Output::from(e) as u32;
@@ -655,33 +644,33 @@ where
655644

656645
let caller = env.caller();
657646

658-
let alpha = pallet_subtensor::Pallet::<T>::do_add_stake(
659-
RawOrigin::Signed(caller.clone()).into(),
660-
hotkey.clone(),
661-
netuid,
662-
tao_amount,
663-
);
647+
let result = transactional::with_transaction(|| {
648+
let alpha = match pallet_subtensor::Pallet::<T>::do_add_stake(
649+
RawOrigin::Signed(caller.clone()).into(),
650+
hotkey.clone(),
651+
netuid,
652+
tao_amount,
653+
) {
654+
Ok(a) => a,
655+
Err(e) => return TransactionOutcome::Rollback(Err(e)),
656+
};
657+
658+
match pallet_subtensor::Pallet::<T>::burn_alpha(
659+
RawOrigin::Signed(caller).into(),
660+
hotkey,
661+
alpha,
662+
netuid,
663+
) {
664+
Ok(_) => TransactionOutcome::Commit(Ok(alpha)),
665+
Err(e) => TransactionOutcome::Rollback(Err(e)),
666+
}
667+
});
664668

665-
match alpha {
669+
match result {
666670
Ok(alpha) => {
667-
let burn_result = pallet_subtensor::Pallet::<T>::burn_alpha(
668-
RawOrigin::Signed(caller).into(),
669-
hotkey,
670-
alpha,
671-
netuid,
672-
);
673-
674-
match burn_result {
675-
Ok(_) => {
676-
env.write_output(&alpha.encode())
677-
.map_err(|_| DispatchError::Other("Failed to write output"))?;
678-
Ok(RetVal::Converging(Output::Success as u32))
679-
}
680-
Err(e) => {
681-
let error_code = Output::from(e) as u32;
682-
Ok(RetVal::Converging(error_code))
683-
}
684-
}
671+
env.write_output(&alpha.encode())
672+
.map_err(|_| DispatchError::Other("Failed to write output"))?;
673+
Ok(RetVal::Converging(Output::Success as u32))
685674
}
686675
Err(e) => {
687676
let error_code = Output::from(e) as u32;

‎chain-extensions/src/tests.rs‎

Lines changed: 142 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1141,8 +1141,8 @@ fn recycle_alpha_clamps_to_available_when_amount_exceeds_stake() {
11411141

11421142
let returned_amount = AlphaBalance::decode(&mut env.output()).unwrap();
11431143
assert_eq!(
1144-
returned_amount, alpha_before,
1145-
"should clamp to available alpha"
1144+
returned_amount, huge_amount,
1145+
"should return requested amount"
11461146
);
11471147

11481148
let alpha_after =
@@ -1242,8 +1242,8 @@ fn burn_alpha_clamps_to_available_when_amount_exceeds_stake() {
12421242

12431243
let returned_amount = AlphaBalance::decode(&mut env.output()).unwrap();
12441244
assert_eq!(
1245-
returned_amount, alpha_before,
1246-
"should clamp to available alpha"
1245+
returned_amount, huge_amount,
1246+
"should return requested amount"
12471247
);
12481248

12491249
let alpha_after =
@@ -1321,6 +1321,144 @@ fn assert_success(ret: RetVal) {
13211321
}
13221322
}
13231323

1324+
#[test]
1325+
fn add_stake_recycle_rollback_on_recycle_failure() {
1326+
mock::new_test_ext(1).execute_with(|| {
1327+
let owner_hotkey = U256::from(12001);
1328+
let owner_coldkey = U256::from(12002);
1329+
let coldkey = U256::from(12101);
1330+
let hotkey = U256::from(12102);
1331+
let min_stake = DefaultMinStake::<mock::Test>::get();
1332+
let tao_amount_raw = min_stake.to_u64().saturating_mul(200);
1333+
1334+
let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey);
1335+
1336+
// Set up very low reserves so recycle will fail with InsufficientLiquidity
1337+
mock::setup_reserves(
1338+
netuid,
1339+
TaoBalance::from(1_000_u64),
1340+
AlphaBalance::from(1_000_u64),
1341+
);
1342+
1343+
mock::register_ok_neuron(netuid, hotkey, coldkey, 0);
1344+
1345+
pallet_subtensor::Pallet::<mock::Test>::add_balance_to_coldkey_account(
1346+
&coldkey,
1347+
TaoBalance::from(tao_amount_raw.saturating_add(1_000_000_000)),
1348+
);
1349+
1350+
let balance_before = pallet_subtensor::Pallet::<mock::Test>::get_coldkey_balance(&coldkey);
1351+
let alpha_before =
1352+
pallet_subtensor::Pallet::<mock::Test>::get_stake_for_hotkey_and_coldkey_on_subnet(
1353+
&hotkey, &coldkey, netuid,
1354+
);
1355+
1356+
let expected_weight = Weight::from_parts(454_200_000, 0)
1357+
.saturating_add(<mock::Test as frame_system::Config>::DbWeight::get().reads(33))
1358+
.saturating_add(<mock::Test as frame_system::Config>::DbWeight::get().writes(19));
1359+
1360+
let mut env = MockEnv::new(
1361+
FunctionId::AddStakeRecycleV1,
1362+
coldkey,
1363+
(hotkey, netuid, TaoBalance::from(tao_amount_raw)).encode(),
1364+
)
1365+
.with_expected_weight(expected_weight);
1366+
1367+
let ret = SubtensorChainExtension::<mock::Test>::dispatch(&mut env).unwrap();
1368+
match ret {
1369+
RetVal::Converging(code) => {
1370+
assert_ne!(code, Output::Success as u32, "should not succeed")
1371+
}
1372+
_ => panic!("unexpected return value"),
1373+
}
1374+
1375+
// Verify full rollback: balance and stake unchanged
1376+
let balance_after = pallet_subtensor::Pallet::<mock::Test>::get_coldkey_balance(&coldkey);
1377+
let alpha_after =
1378+
pallet_subtensor::Pallet::<mock::Test>::get_stake_for_hotkey_and_coldkey_on_subnet(
1379+
&hotkey, &coldkey, netuid,
1380+
);
1381+
1382+
assert_eq!(
1383+
balance_before, balance_after,
1384+
"balance should be unchanged after rollback"
1385+
);
1386+
assert_eq!(
1387+
alpha_before, alpha_after,
1388+
"stake should be unchanged after rollback"
1389+
);
1390+
});
1391+
}
1392+
1393+
#[test]
1394+
fn add_stake_burn_rollback_on_burn_failure() {
1395+
mock::new_test_ext(1).execute_with(|| {
1396+
let owner_hotkey = U256::from(12201);
1397+
let owner_coldkey = U256::from(12202);
1398+
let coldkey = U256::from(12301);
1399+
let hotkey = U256::from(12302);
1400+
let min_stake = DefaultMinStake::<mock::Test>::get();
1401+
let tao_amount_raw = min_stake.to_u64().saturating_mul(200);
1402+
1403+
let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey);
1404+
1405+
// Set up very low reserves so burn will fail with InsufficientLiquidity
1406+
mock::setup_reserves(
1407+
netuid,
1408+
TaoBalance::from(1_000_u64),
1409+
AlphaBalance::from(1_000_u64),
1410+
);
1411+
1412+
mock::register_ok_neuron(netuid, hotkey, coldkey, 0);
1413+
1414+
pallet_subtensor::Pallet::<mock::Test>::add_balance_to_coldkey_account(
1415+
&coldkey,
1416+
TaoBalance::from(tao_amount_raw.saturating_add(1_000_000_000)),
1417+
);
1418+
1419+
let balance_before = pallet_subtensor::Pallet::<mock::Test>::get_coldkey_balance(&coldkey);
1420+
let alpha_before =
1421+
pallet_subtensor::Pallet::<mock::Test>::get_stake_for_hotkey_and_coldkey_on_subnet(
1422+
&hotkey, &coldkey, netuid,
1423+
);
1424+
1425+
let expected_weight = Weight::from_parts(453_000_000, 0)
1426+
.saturating_add(<mock::Test as frame_system::Config>::DbWeight::get().reads(33))
1427+
.saturating_add(<mock::Test as frame_system::Config>::DbWeight::get().writes(18));
1428+
1429+
let mut env = MockEnv::new(
1430+
FunctionId::AddStakeBurnV1,
1431+
coldkey,
1432+
(hotkey, netuid, TaoBalance::from(tao_amount_raw)).encode(),
1433+
)
1434+
.with_expected_weight(expected_weight);
1435+
1436+
let ret = SubtensorChainExtension::<mock::Test>::dispatch(&mut env).unwrap();
1437+
match ret {
1438+
RetVal::Converging(code) => {
1439+
assert_ne!(code, Output::Success as u32, "should not succeed")
1440+
}
1441+
_ => panic!("unexpected return value"),
1442+
}
1443+
1444+
// Verify full rollback: balance and stake unchanged
1445+
let balance_after = pallet_subtensor::Pallet::<mock::Test>::get_coldkey_balance(&coldkey);
1446+
let alpha_after =
1447+
pallet_subtensor::Pallet::<mock::Test>::get_stake_for_hotkey_and_coldkey_on_subnet(
1448+
&hotkey, &coldkey, netuid,
1449+
);
1450+
1451+
assert_eq!(
1452+
balance_before, balance_after,
1453+
"balance should be unchanged after rollback"
1454+
);
1455+
assert_eq!(
1456+
alpha_before, alpha_after,
1457+
"stake should be unchanged after rollback"
1458+
);
1459+
});
1460+
}
1461+
13241462
#[test]
13251463
fn get_stake_info_returns_encoded_runtime_value() {
13261464
mock::new_test_ext(1).execute_with(|| {

0 commit comments

Comments
 (0)