SPARK-3874: Provide stable TaskContext API

This is a small number of clean-up changes on top of #2782. Closes #2782.