Rework cleanup logic

Change-Id: I4c820ce29f1f0fbb4fc0c8243cf9cd71d1a79bbc