Refactor handling of streamConnected

pubsublite/internal/wire/publisher.go

@@ -153,17 +153,16 @@ func (pp *singlePartitionPublisher) onStreamStatusChange(status streamStatus) {
 
 	case streamConnected:
 		pp.unsafeUpdateStatus(serviceActive, nil)
-		pp.enableSendToStream = true
 
 		// To ensure messages are sent in order, we should resend in-flight batches
 		// to the stream immediately after reconnecting, before any new batches.
 		batches := pp.batcher.InFlightBatches()
 		for _, batch := range batches {
 			if !pp.stream.Send(batch.ToPublishRequest()) {
-				pp.enableSendToStream = false
-				break
+				return
 			}
 		}
+		pp.enableSendToStream = true
 
 	case streamTerminated:
 		pp.unsafeInitiateShutdown(serviceTerminated, pp.stream.Error())