325 lines
8.3 KiB
Go
325 lines
8.3 KiB
Go
|
// Copyright 2016 Google Inc. All Rights Reserved.
|
||
|
|
//
|
||
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||
|
|
// you may not use this file except in compliance with the License.
|
||
|
|
// You may obtain a copy of the License at
|
||
|
|
//
|
||
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||
|
|
//
|
||
|
|
// Unless required by applicable law or agreed to in writing, software
|
||
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
|
|
// See the License for the specific language governing permissions and
|
||
|
|
// limitations under the License.
|
||
|
|
|
||
|
|
package pubsub
|
||
|
|
|
||
|
|
import (
|
||
|
|
"fmt"
|
||
|
|
"reflect"
|
||
|
|
"testing"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"golang.org/x/net/context"
|
||
|
|
|
||
|
|
"google.golang.org/api/iterator"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestReturnsDoneOnStop(t *testing.T) {
|
||
|
|
type testCase struct {
|
||
|
|
abort func(*MessageIterator, context.CancelFunc)
|
||
|
|
want error
|
||
|
|
}
|
||
|
|
|
||
|
|
for _, tc := range []testCase{
|
||
|
|
{
|
||
|
|
abort: func(it *MessageIterator, cancel context.CancelFunc) {
|
||
|
|
it.Stop()
|
||
|
|
},
|
||
|
|
want: iterator.Done,
|
||
|
|
},
|
||
|
|
{
|
||
|
|
abort: func(it *MessageIterator, cancel context.CancelFunc) {
|
||
|
|
cancel()
|
||
|
|
},
|
||
|
|
want: context.Canceled,
|
||
|
|
},
|
||
|
|
{
|
||
|
|
abort: func(it *MessageIterator, cancel context.CancelFunc) {
|
||
|
|
it.Stop()
|
||
|
|
cancel()
|
||
|
|
},
|
||
|
|
want: iterator.Done,
|
||
|
|
},
|
||
|
|
{
|
||
|
|
abort: func(it *MessageIterator, cancel context.CancelFunc) {
|
||
|
|
cancel()
|
||
|
|
it.Stop()
|
||
|
|
},
|
||
|
|
want: iterator.Done,
|
||
|
|
},
|
||
|
|
} {
|
||
|
|
s := &blockingFetch{}
|
||
|
|
ctx, cancel := context.WithCancel(context.Background())
|
||
|
|
it := newMessageIterator(ctx, s, "subname", &pullOptions{ackDeadline: time.Second * 10, maxExtension: time.Hour})
|
||
|
|
defer it.Stop()
|
||
|
|
tc.abort(it, cancel)
|
||
|
|
|
||
|
|
_, err := it.Next()
|
||
|
|
if err != tc.want {
|
||
|
|
t.Errorf("iterator Next error after abort: got:\n%v\nwant:\n%v", err, tc.want)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// blockingFetch implements message fetching by not returning until its context is cancelled.
|
||
|
|
type blockingFetch struct {
|
||
|
|
service
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *blockingFetch) fetchMessages(ctx context.Context, subName string, maxMessages int32) ([]*Message, error) {
|
||
|
|
<-ctx.Done()
|
||
|
|
return nil, ctx.Err()
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *blockingFetch) newStreamingPuller(ctx context.Context, subName string, ackDeadline int32) *streamingPuller {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// justInTimeFetch simulates the situation where the iterator is aborted just after the fetch RPC
|
||
|
|
// succeeds, so the rest of puller.Next will continue to execute and return sucessfully.
|
||
|
|
type justInTimeFetch struct {
|
||
|
|
service
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *justInTimeFetch) fetchMessages(ctx context.Context, subName string, maxMessages int32) ([]*Message, error) {
|
||
|
|
<-ctx.Done()
|
||
|
|
// The context was cancelled, but let's pretend that this happend just after our RPC returned.
|
||
|
|
|
||
|
|
var result []*Message
|
||
|
|
for i := 0; i < int(maxMessages); i++ {
|
||
|
|
val := fmt.Sprintf("msg%v", i)
|
||
|
|
result = append(result, &Message{Data: []byte(val), ackID: val})
|
||
|
|
}
|
||
|
|
return result, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *justInTimeFetch) splitAckIDs(ids []string) ([]string, []string) {
|
||
|
|
return nil, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *justInTimeFetch) modifyAckDeadline(ctx context.Context, subName string, deadline time.Duration, ackIDs []string) error {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *justInTimeFetch) newStreamingPuller(ctx context.Context, subName string, ackDeadline int32) *streamingPuller {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestAfterAbortReturnsNoMoreThanOneMessage(t *testing.T) {
|
||
|
|
// Each test case is excercised by making two concurrent blocking calls on a
|
||
|
|
// MessageIterator, and then aborting the iterator.
|
||
|
|
// The result should be one call to Next returning a message, and the other returning an error.
|
||
|
|
type testCase struct {
|
||
|
|
abort func(*MessageIterator, context.CancelFunc)
|
||
|
|
// want is the error that should be returned from one Next invocation.
|
||
|
|
want error
|
||
|
|
}
|
||
|
|
for n := 1; n < 3; n++ {
|
||
|
|
for _, tc := range []testCase{
|
||
|
|
{
|
||
|
|
abort: func(it *MessageIterator, cancel context.CancelFunc) {
|
||
|
|
it.Stop()
|
||
|
|
},
|
||
|
|
want: iterator.Done,
|
||
|
|
},
|
||
|
|
{
|
||
|
|
abort: func(it *MessageIterator, cancel context.CancelFunc) {
|
||
|
|
cancel()
|
||
|
|
},
|
||
|
|
want: context.Canceled,
|
||
|
|
},
|
||
|
|
{
|
||
|
|
abort: func(it *MessageIterator, cancel context.CancelFunc) {
|
||
|
|
it.Stop()
|
||
|
|
cancel()
|
||
|
|
},
|
||
|
|
want: iterator.Done,
|
||
|
|
},
|
||
|
|
{
|
||
|
|
abort: func(it *MessageIterator, cancel context.CancelFunc) {
|
||
|
|
cancel()
|
||
|
|
it.Stop()
|
||
|
|
},
|
||
|
|
want: iterator.Done,
|
||
|
|
},
|
||
|
|
} {
|
||
|
|
s := &justInTimeFetch{}
|
||
|
|
ctx, cancel := context.WithCancel(context.Background())
|
||
|
|
|
||
|
|
// if maxPrefetch == 1, there will be no messages in the puller buffer when Next is invoked the second time.
|
||
|
|
// if maxPrefetch == 2, there will be 1 message in the puller buffer when Next is invoked the second time.
|
||
|
|
po := &pullOptions{
|
||
|
|
ackDeadline: time.Second * 10,
|
||
|
|
maxExtension: time.Hour,
|
||
|
|
maxPrefetch: int32(n),
|
||
|
|
}
|
||
|
|
it := newMessageIterator(ctx, s, "subname", po)
|
||
|
|
defer it.Stop()
|
||
|
|
|
||
|
|
type result struct {
|
||
|
|
m *Message
|
||
|
|
err error
|
||
|
|
}
|
||
|
|
results := make(chan *result, 2)
|
||
|
|
|
||
|
|
for i := 0; i < 2; i++ {
|
||
|
|
go func() {
|
||
|
|
m, err := it.Next()
|
||
|
|
results <- &result{m, err}
|
||
|
|
if err == nil {
|
||
|
|
m.Done(false)
|
||
|
|
}
|
||
|
|
}()
|
||
|
|
}
|
||
|
|
// Wait for goroutines to block on it.Next().
|
||
|
|
time.Sleep(time.Millisecond)
|
||
|
|
tc.abort(it, cancel)
|
||
|
|
|
||
|
|
result1 := <-results
|
||
|
|
result2 := <-results
|
||
|
|
|
||
|
|
// There should be one error result, and one non-error result.
|
||
|
|
// Make result1 be the non-error result.
|
||
|
|
if result1.err != nil {
|
||
|
|
result1, result2 = result2, result1
|
||
|
|
}
|
||
|
|
|
||
|
|
if string(result1.m.Data) != "msg0" {
|
||
|
|
t.Errorf("After abort, got message: %v, want %v", result1.m.Data, "msg0")
|
||
|
|
}
|
||
|
|
if result1.err != nil {
|
||
|
|
t.Errorf("After abort, got : %v, want nil", result1.err)
|
||
|
|
}
|
||
|
|
if result2.m != nil {
|
||
|
|
t.Errorf("After abort, got message: %v, want nil", result2.m)
|
||
|
|
}
|
||
|
|
if result2.err != tc.want {
|
||
|
|
t.Errorf("After abort, got err: %v, want %v", result2.err, tc.want)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
type fetcherServiceWithModifyAckDeadline struct {
|
||
|
|
fetcherService
|
||
|
|
events chan string
|
||
|
|
}
|
||
|
|
|
||
|
|
func (f *fetcherServiceWithModifyAckDeadline) modifyAckDeadline(_ context.Context, _ string, d time.Duration, ids []string) error {
|
||
|
|
// Different versions of Go use different representations for time.Duration(0).
|
||
|
|
var ds string
|
||
|
|
if d == 0 {
|
||
|
|
ds = "0s"
|
||
|
|
} else {
|
||
|
|
ds = d.String()
|
||
|
|
}
|
||
|
|
f.events <- fmt.Sprintf("modAck(%v, %s)", ids, ds)
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (f *fetcherServiceWithModifyAckDeadline) splitAckIDs(ackIDs []string) ([]string, []string) {
|
||
|
|
return ackIDs, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (f *fetcherServiceWithModifyAckDeadline) newStreamingPuller(ctx context.Context, subName string, ackDeadline int32) *streamingPuller {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestMultipleStopCallsBlockUntilMessageDone(t *testing.T) {
|
||
|
|
events := make(chan string, 3)
|
||
|
|
s := &fetcherServiceWithModifyAckDeadline{
|
||
|
|
fetcherService{
|
||
|
|
results: []fetchResult{
|
||
|
|
{
|
||
|
|
msgs: []*Message{{ackID: "a"}, {ackID: "b"}},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
events,
|
||
|
|
}
|
||
|
|
|
||
|
|
ctx := context.Background()
|
||
|
|
it := newMessageIterator(ctx, s, "subname", &pullOptions{ackDeadline: time.Second * 10, maxExtension: 0})
|
||
|
|
|
||
|
|
m, err := it.Next()
|
||
|
|
if err != nil {
|
||
|
|
t.Errorf("error calling Next: %v", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
go func() {
|
||
|
|
it.Stop()
|
||
|
|
events <- "stopped"
|
||
|
|
}()
|
||
|
|
go func() {
|
||
|
|
it.Stop()
|
||
|
|
events <- "stopped"
|
||
|
|
}()
|
||
|
|
|
||
|
|
time.Sleep(10 * time.Millisecond)
|
||
|
|
m.Done(false)
|
||
|
|
|
||
|
|
got := []string{<-events, <-events, <-events}
|
||
|
|
want := []string{"modAck([a], 0s)", "stopped", "stopped"}
|
||
|
|
if !reflect.DeepEqual(got, want) {
|
||
|
|
t.Errorf("stopping iterator, got: %v ; want: %v", got, want)
|
||
|
|
}
|
||
|
|
|
||
|
|
// The iterator is stopped, so should not return another message.
|
||
|
|
m, err = it.Next()
|
||
|
|
if m != nil {
|
||
|
|
t.Errorf("message got: %v ; want: nil", m)
|
||
|
|
}
|
||
|
|
if err != iterator.Done {
|
||
|
|
t.Errorf("err got: %v ; want: %v", err, iterator.Done)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestFastNack(t *testing.T) {
|
||
|
|
events := make(chan string, 3)
|
||
|
|
s := &fetcherServiceWithModifyAckDeadline{
|
||
|
|
fetcherService{
|
||
|
|
results: []fetchResult{
|
||
|
|
{
|
||
|
|
msgs: []*Message{{ackID: "a"}, {ackID: "b"}},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
events,
|
||
|
|
}
|
||
|
|
|
||
|
|
ctx := context.Background()
|
||
|
|
it := newMessageIterator(ctx, s, "subname", &pullOptions{
|
||
|
|
ackDeadline: time.Second * 6,
|
||
|
|
maxExtension: time.Second * 10,
|
||
|
|
})
|
||
|
|
// Get both messages.
|
||
|
|
_, err := it.Next()
|
||
|
|
if err != nil {
|
||
|
|
t.Errorf("error calling Next: %v", err)
|
||
|
|
}
|
||
|
|
m2, err := it.Next()
|
||
|
|
if err != nil {
|
||
|
|
t.Errorf("error calling Next: %v", err)
|
||
|
|
}
|
||
|
|
// Ignore the first, nack the second.
|
||
|
|
m2.Done(false)
|
||
|
|
|
||
|
|
got := []string{<-events, <-events}
|
||
|
|
// The nack should happen before the deadline extension.
|
||
|
|
want := []string{"modAck([b], 0s)", "modAck([a], 6s)"}
|
||
|
|
if !reflect.DeepEqual(got, want) {
|
||
|
|
t.Errorf("got: %v ; want: %v", got, want)
|
||
|
|
}
|
||
|
|
}
|