fix(pubsub): prevent draining error return for Receive (#4733)

pubsub/internal/scheduler/receive_scheduler.go

@@ -15,7 +15,6 @@
 package scheduler
 
 import (
-	"errors"
 	"sync"
 )
 
@@ -68,12 +67,6 @@ func NewReceiveScheduler(workers int) *ReceiveScheduler {
 // call causes pushback to the pubsub service (less Receive calls on the
 // long-lived stream), which keeps memory footprint stable.
 func (s *ReceiveScheduler) Add(key string, item interface{}, handle func(item interface{})) error {
-	select {
-	case <-s.done:
-		return errors.New("draining")
-	default:
-	}
-
 	if key == "" {
 		// Spawn a worker.
 		s.workers <- struct{}{}

pubsub/subscription.go

@@ -925,13 +925,29 @@ func (s *Subscription) Receive(ctx context.Context, f func(context.Context, *Mes
 						}
 					}
 				}
+				// If the context is done, don't pull more messages.
+				select {
+				case <-ctx.Done():
+					return nil
+				default:
+				}
 				msgs, err := iter.receive(maxToPull)
 				if err == io.EOF {
 					return nil
 				}
 				if err != nil {
 					return err
 				}
+				// If context is done and messages have been pulled,
+				// nack them.
+				select {
+				case <-ctx.Done():
+					for _, m := range msgs {
+						m.Nack()
+					}
+					return nil
+				default:
+				}
 				for i, msg := range msgs {
 					msg := msg
 					// TODO(jba): call acquire closer to when the message is allocated.