Refactor TrafficGenerator