Function TestBatchSizeDecays

pubsub/pubsub_test.go:387–445

TestBatchSizeDecay verifies that the batch size decays when no messages are available. (see https://github.com/google/go-cloud/issues/2849).

(t *testing.T)

// TestBatchSizeDecay verifies that the batch size decays when no messages are available.
// (see https://github.com/google/go-cloud/issues/2849).
func TestBatchSizeDecays(t *testing.T) {
	ctx := context.Background()
	fs := &failSub{}
	// Allow multiple handlers and cap max batch size to ensure we get concurrency.
	sub := NewSubscription(fs, &batcher.Options{MaxHandlers: 10, MaxBatchSize: 2}, nil)
	defer sub.Shutdown(ctx)

	// Records the last batch size.
	var mu sync.Mutex
	lastMaxMessages := 0
	sub.preReceiveBatchHook = func(maxMessages int) {
		mu.Lock()
		defer mu.Unlock()
		lastMaxMessages = maxMessages
	}

	// Do some receives to allow the number of batches to increase past 1.
	for range 100 {
		m, err := sub.Receive(ctx)
		if err != nil {
			t.Fatalf("Receive: got %v, want nil", err)
		}
		m.Ack()
	}

	// Tell the failSub to start returning no messages.
	fs.mu.Lock()
	fs.empty = true
	fs.mu.Unlock()

	mu.Lock()
	highWaterMarkBatchSize := lastMaxMessages
	if lastMaxMessages <= 1 {
		t.Fatal("max messages wasn't greater than 1")
	}
	mu.Unlock()

	// Make a bunch of calls to Receive to drain any outstanding
	// messages, and wait some extra time during which we should
	// continue polling, and the batch size should decay.
	for {
		ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
		defer cancel()
		m, err := sub.Receive(ctx)
		if err != nil {
			// Expected: no more messages, and timed out.
			break
		}
		// Drained a message.
		m.Ack()
	}

	// Verify that the batch size decayed.
	mu.Lock()
	if lastMaxMessages >= highWaterMarkBatchSize {
		t.Fatalf("wanted batch size to decay; high water mark was %d, now %d", highWaterMarkBatchSize, lastMaxMessages)
	}
	mu.Unlock()