[SPARK-3543] TaskContext remaining cleanup work.