some code cleanup. make restart more smarter