Cleanup code a little bit, nicer design, fixed cancel handling