From fc08c49fc013cbad00642bbba317e02f0ba15a6d Mon Sep 17 00:00:00 2001
From: Alex Hong <9397363+hongalex@users.noreply.github.com>
Date: Sun, 23 Jan 2022 23:53:23 -0800
Subject: [PATCH] fix(pubsub): pass context into checkOrdering to allow cancel
 (#5316)

---
 pubsub/subscription.go      | 5 ++---
 pubsub/subscription_test.go | 7 +++++++
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/pubsub/subscription.go b/pubsub/subscription.go
index 1a741a78e5a..b3bb646a41e 100644
--- a/pubsub/subscription.go
+++ b/pubsub/subscription.go
@@ -840,7 +840,7 @@ func (s *Subscription) Receive(ctx context.Context, f func(context.Context, *Mes
 	s.mu.Unlock()
 	defer func() { s.mu.Lock(); s.receiveActive = false; s.mu.Unlock() }()
 
-	s.checkOrdering()
+	s.checkOrdering(ctx)
 
 	maxCount := s.ReceiveSettings.MaxOutstandingMessages
 	if maxCount == 0 {
@@ -1033,8 +1033,7 @@ func (s *Subscription) Receive(ctx context.Context, f func(context.Context, *Mes
 // the roles/viewer or roles/pubsub.viewer role) we will assume
 // EnableMessageOrdering to be true.
 // See: https://github.com/googleapis/google-cloud-go/issues/3884
-func (s *Subscription) checkOrdering() {
-	ctx := context.Background()
+func (s *Subscription) checkOrdering(ctx context.Context) {
 	cfg, err := s.Config(ctx)
 	if err != nil {
 		s.enableOrdering = true
diff --git a/pubsub/subscription_test.go b/pubsub/subscription_test.go
index 4f35457733b..7595905152b 100644
--- a/pubsub/subscription_test.go
+++ b/pubsub/subscription_test.go
@@ -427,4 +427,11 @@ func TestOrdering_CreateSubscription(t *testing.T) {
 	if !cfg.EnableMessageOrdering {
 		t.Fatalf("Expected EnableMessageOrdering to be true in %s", orderSub.String())
 	}
+
+	// Test cancellation works as intended with ordering enabled.
+	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
+	defer cancel()
+	orderSub.Receive(ctx, func(ctx context.Context, msg *Message) {
+		msg.Ack()
+	})
 }