Added shutdown hook
Refactoring