Improve partition management for kafka 0.8