Improve data partition waiting code