Improve shutdown function

fixes #10 and #11