In order to support true high availability TaskManager and JobManager have to be able to distinguish whether a message was sent from the leader or whether a message was sent from a former leader. Messages which come from a former leader have to be discarded in order to guarantee a consistent state.
A way to do achieve this is to assign a leader session ID to a JobManager once he's elected as leader. This leader session ID is sent to the TaskManager upon registration at the JobManager. All subsequent messages should then be decorated with this leader session ID to mark them as valid. On the TaskManager side the received leader session ID as a response to the registration message, can then be used to validate incoming messages.
The same holds true for registration messages which should have a registration session ID, too. That way, it is possible to distinguish invalid registration messages from valid ones. The registration session ID can be assigned once the TaskManager is notified about the new leader.