Add some cleanup