[SPARK-4088] [PySpark] Python worker should exit after socket is closed by JVM

In case of take() or exception in Python, python worker may exit before JVM read() all the response, then the write thread may raise "Connection reset" exception.

Python should always wait JVM to close the socket first.

cc @JoshRosen This is a warm fix, or the tests will be flaky.