diff --git a/src/EjectionManager.sol b/src/EjectionManager.sol index a9a4bc77..cbb595cb 100644 --- a/src/EjectionManager.sol +++ b/src/EjectionManager.sol @@ -60,17 +60,21 @@ contract EjectionManager is IEjectionManager, OwnableUpgradeable{ /** * @notice Ejects operators from the AVSs RegistryCoordinator under a ratelimit * @param _operatorIds The ids of the operators 'j' to eject for each quorum 'i' + * @return ejectedOperatorsForQuorum The total number of operators ejected for each quorum 'i' * @dev This function will eject as many operators as possible without reverting prioritizing operators at the lower index * @dev The owner can eject operators without recording of stake ejection */ - function ejectOperators(bytes32[][] memory _operatorIds) external { + function ejectOperators(bytes32[][] memory _operatorIds) external returns (uint256[] memory){ require(isEjector[msg.sender] || msg.sender == owner(), "Ejector: Only owner or ejector can eject"); + uint256[] memory ejectedOperatorsForQuorum = new uint256[](_operatorIds.length); + for(uint i = 0; i < _operatorIds.length; ++i) { uint8 quorumNumber = uint8(i); uint256 amountEjectable = amountEjectableForQuorum(quorumNumber); - uint256 stakeForEjection; + uint256 stakeForEjection; + uint256 ejectedOperators; bool broke; for(uint8 j = 0; j < _operatorIds[i].length; ++j) { @@ -96,6 +100,7 @@ contract EjectionManager is IEjectionManager, OwnableUpgradeable{ abi.encodePacked(quorumNumber) ) { stakeForEjection += operatorStake; + ++ejectedOperators; emit OperatorEjected(_operatorIds[i][j], quorumNumber); } catch (bytes memory err) { emit FailedOperatorEjection(_operatorIds[i][j], quorumNumber, err); @@ -110,7 +115,10 @@ contract EjectionManager is IEjectionManager, OwnableUpgradeable{ })); } + ejectedOperatorsForQuorum[i] = ejectedOperators; } + + return ejectedOperatorsForQuorum; } /** diff --git a/src/interfaces/IEjectionManager.sol b/src/interfaces/IEjectionManager.sol index 83299484..6cebcddc 100644 --- a/src/interfaces/IEjectionManager.sol +++ b/src/interfaces/IEjectionManager.sol @@ -32,7 +32,7 @@ interface IEjectionManager { * @notice Ejects operators from the AVSs registryCoordinator under a ratelimit * @param _operatorIds The ids of the operators to eject for each quorum */ - function ejectOperators(bytes32[][] memory _operatorIds) external; + function ejectOperators(bytes32[][] memory _operatorIds) external returns (uint256[] memory); /** * @notice Sets the ratelimit parameters for a quorum diff --git a/test/unit/EjectionManagerUnit.t.sol b/test/unit/EjectionManagerUnit.t.sol index 82a22daa..9bdddc88 100644 --- a/test/unit/EjectionManagerUnit.t.sol +++ b/test/unit/EjectionManagerUnit.t.sol @@ -87,7 +87,11 @@ contract EjectionManagerUnitTests is MockAVSDeployer { } cheats.prank(ejector); - ejectionManager.ejectOperators(operatorIds); + uint256[] memory ejectedOperatorsForQuorum = ejectionManager.ejectOperators(operatorIds); + + for(uint8 i = 0; i < numQuorums; i++) { + assertEq(ejectedOperatorsForQuorum[i], operatorsToEject); + } assertEq(uint8(registryCoordinator.getOperatorStatus(defaultOperator)), uint8(IRegistryCoordinator.OperatorStatus.DEREGISTERED)); } @@ -118,7 +122,11 @@ contract EjectionManagerUnitTests is MockAVSDeployer { } cheats.prank(ejector); - ejectionManager.ejectOperators(operatorIds); + uint256[] memory ejectedOperatorsForQuorum = ejectionManager.ejectOperators(operatorIds); + + for(uint8 i = 0; i < numQuorums; i++) { + assertEq(ejectedOperatorsForQuorum[i], operatorsToEject); + } for(uint8 i = 0; i < operatorsToEject; i++) { assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.DEREGISTERED)); @@ -152,7 +160,11 @@ contract EjectionManagerUnitTests is MockAVSDeployer { } cheats.prank(ejector); - ejectionManager.ejectOperators(operatorIds); + uint256[] memory ejectedOperatorsForQuorum = ejectionManager.ejectOperators(operatorIds); + + for(uint8 i = 0; i < numQuorums; i++) { + assertEq(ejectedOperatorsForQuorum[i], operatorsCanEject); + } for(uint8 i = 0; i < operatorsCanEject; i++) { assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.DEREGISTERED)); @@ -189,7 +201,11 @@ contract EjectionManagerUnitTests is MockAVSDeployer { } cheats.prank(ejector); - ejectionManager.ejectOperators(operatorIds); + uint256[] memory ejectedOperatorsForQuorum = ejectionManager.ejectOperators(operatorIds); + + for(uint8 i = 0; i < numQuorums; i++) { + assertEq(ejectedOperatorsForQuorum[i], operatorsToEject); + } for(uint8 i = 0; i < operatorsToEject; i++) { assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.DEREGISTERED)); @@ -217,7 +233,11 @@ contract EjectionManagerUnitTests is MockAVSDeployer { } cheats.prank(ejector); - ejectionManager.ejectOperators(operatorIds); + ejectedOperatorsForQuorum = ejectionManager.ejectOperators(operatorIds); + + for(uint8 i = 0; i < numQuorums; i++) { + assertEq(ejectedOperatorsForQuorum[i], operatorsToEject); + } for(uint8 i = 0; i < operatorsToEject; i++) { assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, operatorsToEject + i))), uint8(IRegistryCoordinator.OperatorStatus.DEREGISTERED)); @@ -255,7 +275,11 @@ contract EjectionManagerUnitTests is MockAVSDeployer { } cheats.prank(ejector); - ejectionManager.ejectOperators(operatorIds); + uint256[] memory ejectedOperatorsForQuorum = ejectionManager.ejectOperators(operatorIds); + + for(uint8 i = 0; i < numQuorums; i++) { + assertEq(ejectedOperatorsForQuorum[i], operatorsToEject); + } for(uint8 i = 0; i < operatorsToEject; i++) { assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.DEREGISTERED)); @@ -288,7 +312,11 @@ contract EjectionManagerUnitTests is MockAVSDeployer { } cheats.prank(registryCoordinatorOwner); - ejectionManager.ejectOperators(operatorIds); + uint256[] memory ejectedOperatorsForQuorum = ejectionManager.ejectOperators(operatorIds); + + for(uint8 i = 0; i < numQuorums; i++) { + assertEq(ejectedOperatorsForQuorum[i], operatorsToEject); + } for(uint8 i = 0; i < operatorsToEject; i++) { assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.DEREGISTERED)); @@ -326,7 +354,11 @@ contract EjectionManagerUnitTests is MockAVSDeployer { } cheats.prank(ejector); - ejectionManager.ejectOperators(operatorIds); + uint256[] memory ejectedOperatorsForQuorum = ejectionManager.ejectOperators(operatorIds); + + for(uint8 i = 0; i < numQuorums; i++) { + assertEq(ejectedOperatorsForQuorum[i], operatorsToEject - 1); + } for(uint8 i = 0; i < operatorsToEject; i++) { assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.DEREGISTERED));