Skip to content

Commit

Permalink
fix(pubsub): Closes googleapis#10094 - memory leak in pubsub receive (g…
Browse files Browse the repository at this point in the history
…oogleapis#10153)

* fix(pubsub): make sure grpc stream gets closed

* fix(pubsub): preserve cancellation error behaviour

* fix(pubsub): the stream cancellation error is not necessarily a grpc error

---------

Co-authored-by: Alex Hong <9397363+hongalex@users.noreply.github.com>
  • Loading branch information
itizir and hongalex committed May 17, 2024
1 parent b371210 commit 66581c4
Showing 1 changed file with 11 additions and 0 deletions.
11 changes: 11 additions & 0 deletions pubsub/iterator.go
Original file line number Diff line number Diff line change
Expand Up @@ -163,6 +163,9 @@ func (it *messageIterator) stop() {
it.checkDrained()
it.mu.Unlock()
it.wg.Wait()
if it.ps != nil {
it.ps.cancel()
}
}

// checkDrained closes the drained channel if the iterator has been stopped and all
Expand Down Expand Up @@ -246,6 +249,14 @@ func (it *messageIterator) receive(maxToPull int32) ([]*Message, error) {
rmsgs, err = it.pullMessages(maxToPull)
} else {
rmsgs, err = it.recvMessages()
// If stopping the iterator results in the grpc stream getting shut down and
// returning an error here, treat the same as above and return EOF.
// If the cancellation comes from the underlying grpc client getting closed,
// do propagate the cancellation error.
// See https://github.com/googleapis/google-cloud-go/pull/10153#discussion_r1600814775
if err != nil && it.ps.ctx.Err() == context.Canceled {
err = io.EOF
}
}
// Any error here is fatal.
if err != nil {
Expand Down

0 comments on commit 66581c4

Please sign in to comment.