What
Make improvements as suggested in this pr

Why
We merged that pr to unblock another one but there was some feedback items to improve it.