Refactor and fix lock-up
Merge pull request #242 from my-personal-forks/refactor

Refactor and fix lock-up