fix: Handle TransportError Exceptions thrown from gapic_publish

[mukund-ananthu]

Current behavior:

If TransportError is thrown by google.auth library:

https://github.com/googleapis/google-auth-library-python/blob/c3ea09fd8b9ee8f094263fdc809c59d9dfaa3124/google/auth/compute_engine/_metadata.py#L187-L189

when doing a _gapic_publish:

python-pubsub/google/cloud/pubsub_v1/publisher/_batch/thread.py

Lines 322 to 327 in f79d35a

	response = self._client._gapic_publish(
	topic=self._topic,
	messages=[wrapper.message for wrapper in self._message_wrappers],
	retry=self._commit_retry,
	timeout=self._commit_timeout,
	)

Then,

1. The exception is not caught:

python-pubsub/google/cloud/pubsub_v1/publisher/_batch/thread.py

Line 345 in f79d35a

except google.api_core.exceptions.GoogleAPIError as exc:

2. Status of batch not set to ERROR:

python-pubsub/google/cloud/pubsub_v1/publisher/_batch/thread.py

Line 348 in f79d35a

self._status = base.BatchStatus.ERROR

3. Exceptions not set in the publish futures

python-pubsub/google/cloud/pubsub_v1/publisher/_batch/thread.py

Lines 364 to 369 in f79d35a

	if self._batch_done_callback is not None:
	# Failed to publish batch.
	self._batch_done_callback(batch_transport_succeeded)
	for future in self._futures:
	future.set_exception(exc)

4. The thread in which the batch gapic_publish was occuring dies without setting the correct batch state / future exceptions:

python-pubsub/google/cloud/pubsub_v1/publisher/_batch/thread.py

Lines 237 to 239 in f79d35a

	commit_thread = threading.Thread(
	name="Thread-CommitBatchPublisher", target=self._commit, daemon=True
	)

5. State of Batch remains IN_PROGRESS indefinitely

python-pubsub/google/cloud/pubsub_v1/publisher/_batch/thread.py

Line 288 in f79d35a

self._status = base.BatchStatus.IN_PROGRESS

6. Subsequent publishes of messages result in creation of a new batch:

python-pubsub/google/cloud/pubsub_v1/publisher/_batch/thread.py

Lines 449 to 450 in f79d35a

	if self.status != base.BatchStatus.ACCEPTING_MESSAGES:
	return None

python-pubsub/google/cloud/pubsub_v1/publisher/_sequencer/ordered_sequencer.py

Lines 323 to 326 in f79d35a

	while future is None:
	batch = self._create_batch(commit_retry=retry, commit_timeout=timeout)
	self._ordered_batches.append(batch)
	future = batch.publish(wrapper)

7. Messages that were in the batch whose commit thread died:

• Will not have the TransportError set as an exception in their publish futures
• The publish future will either hang or timeout

Behavior expected with change

1. TransportError during gapic_publish will be caught in addition to the existing GoogleAPIError

python-pubsub/google/cloud/pubsub_v1/publisher/_batch/thread.py

Line 345 in f79d35a

except google.api_core.exceptions.GoogleAPIError as exc:

2. Existing behavior when GoogleAPIError is caught during gapic_publish would also now apply to TransportError, i.e:

• Batch status will be set to ERROR, future exceptions will be set:

python-pubsub/google/cloud/pubsub_v1/publisher/_batch/thread.py

Lines 345 to 371 in f79d35a

	except google.api_core.exceptions.GoogleAPIError as exc:
	# We failed to publish, even after retries, so set the exception on
	# all futures and exit.
	self._status = base.BatchStatus.ERROR
	if self._client.open_telemetry_enabled:
	if self._rpc_span:
	self._rpc_span.record_exception(
	exception=exc,
	)
	self._rpc_span.set_status(
	trace.Status(status_code=trace.StatusCode.ERROR)
	)
	self._rpc_span.end()
	for wrapper in self._message_wrappers:
	wrapper.end_create_span(exc=exc)
	batch_transport_succeeded = False
	if self._batch_done_callback is not None:
	# Failed to publish batch.
	self._batch_done_callback(batch_transport_succeeded)
	for future in self._futures:
	future.set_exception(exc)
	return

• Subsequent publish of messages on the batch object will cause an AssertionError:

python-pubsub/google/cloud/pubsub_v1/publisher/_batch/thread.py

Lines 444 to 447 in f79d35a

	with self._state_lock:
	assert (
	self._status != base.BatchStatus.ERROR
	), "Publish after stop() or publish error."

which will be bubbled up through the Ordered / Unordered Sequenceer's publish method:

python-pubsub/google/cloud/pubsub_v1/publisher/_sequencer/ordered_sequencer.py

Lines 321 to 328 in f79d35a

	batch = self._ordered_batches[-1]
	future = batch.publish(wrapper)
	while future is None:
	batch = self._create_batch(commit_retry=retry, commit_timeout=timeout)
	self._ordered_batches.append(batch)
	future = batch.publish(wrapper)
	return future

to the Publisher client's publish method:

python-pubsub/google/cloud/pubsub_v1/publisher/client.py

Lines 477 to 495 in f79d35a

	future = sequencer.publish(
	wrapper=wrapper, retry=retry, timeout=timeout
	)
	future.add_done_callback(on_publish_done)
	except BaseException as be:
	# Exceptions can be thrown when attempting to add messages to
	# the batch. If they're thrown, record them in publisher
	# batching and create span, end the spans and bubble the
	# exception up.
	if self._open_telemetry_enabled:
	if wrapper:
	wrapper.end_publisher_batching_span(be)
	wrapper.end_create_span(be)
	else: # pragma: NO COVER
	warnings.warn(
	message="PublishMessageWrapper is None. Hence, not recording exception and ending publisher batching span and create span",
	category=RuntimeWarning,
	)
	raise be

which would bubble the exception back to the user

Additional details can be found here: #1173 (comment)

Fixes #1173 🦕