Refactor cleanup to own function
Improve docs