Fix and simplify rollback system
Add test for rollback