Fix for hang on shutdown

There's some cleanups in here too. Which should be okay because the actual fix is a one line change.
Reviewers: @nbougalis @ximinez