Error handling change and small cleanups