LogDog: Add collector service implementation.

common/gcloud/gcps/subscriber/subscriber_test.go

 // Copyright 2015 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 package subscriber
 
 import (
         "fmt"
         "sort"
         "sync"
+        "sync/atomic"
         "testing"
         "time"
 
         "github.com/luci/luci-go/common/clock"
         "github.com/luci/luci-go/common/clock/testclock"
         "github.com/luci/luci-go/common/errors"
         "github.com/luci/luci-go/common/gcloud/gcps"
         "golang.org/x/net/context"
         "google.golang.org/cloud/pubsub"
 
+        . "github.com/luci/luci-go/common/testing/assertions"
         . "github.com/smartystreets/goconvey/convey"
 )
 
-type testPubSub struct {
-        c     context.Context
-        errCB func() error
-        sub   gcps.Subscription
-        msgC  chan *pubsub.Message
+type event struct {
+        msg *pubsub.Message
+        err error
 }
 
-func (ps *testPubSub) Pull(s gcps.Subscription, amount int) ([]*pubsub.Message, error) {
-        if ps.errCB != nil {
-                if err := ps.errCB(); err != nil {
-                        return nil, err
-                }
-        }
-        if ps.sub != s {
-                return nil, fmt.Errorf("unknown subscription %q", s)
-        }
-        if amount <= 0 {
-                return nil, fmt.Errorf("invalid amount: %d", amount)
-        }
-
-        select {
-        case <-ps.c.Done():
-                return nil, ps.c.Err()
-
-        case msg := <-ps.msgC:
-                return []*pubsub.Message{msg}, nil
-        }
+type testSource struct {
+        sub    gcps.Subscription
+        eventC chan event
 }
 
-func (ps *testPubSub) send(s ...string) {
-        for _, v := range s {
-                ps.msgC <- &pubsub.Message{
-                        ID:   v,
-                        Data: []byte(v),
+func (s *testSource) Pull(c context.Context) ([]*pubsub.Message, error) {
+        select {
+        case <-c.Done():
+                return nil, c.Err()
+
+        case e := <-s.eventC:
+                switch {
+                case e.err != nil:
+                        return nil, e.err
+
+                case e.msg != nil:
+                        return []*pubsub.Message{e.msg}, nil
+
+                default:
+                        return nil, nil
                 }
         }
 }
 
+func (s *testSource) message(id ...string) {
+        for _, v := range id {
+                if v != "" {
+                        s.eventC <- event{msg: &pubsub.Message{
+                                ID:    v,
+                                AckID: v,
+                                Data:  []byte(v),
+                        }}
+                } else {
+                        s.eventC <- event{}
+                }
+        }
+}
+
+func (s *testSource) error(err error) {
+        s.eventC <- event{err: err}
+}
+
+type testACK struct {
+        sync.Mutex
+
+        acks map[string]struct{}
+}
+
+func (a *testACK) Ack(id string) {
+        a.Lock()
+        defer a.Unlock()
+
+        if a.acks == nil {
+                a.acks = make(map[string]struct{})
+        }
+        a.acks[id] = struct{}{}
+}
+
+func (a *testACK) getACKs() []string {
+        a.Lock()
+        defer a.Unlock()
+        return dumpStringSet(a.acks)
+}
+
+func dumpStringSet(s map[string]struct{}) []string {
+        v := make([]string, 0, len(s))
+        for a := range s {
+                v = append(v, a)
+        }
+        sort.Strings(v)
+        return v
+}
+
 func TestSubscriber(t *testing.T) {
         t.Parallel()
 
         Convey(`A Subscriber configuration using a testing Pub/Sub`, t, func() {
                 c := context.Background()
                 c, tc := testclock.UseTime(c, testclock.TestTimeLocal)
 
                 c, cancelFunc := context.WithCancel(c)
                 defer cancelFunc()
 
-                ps := &testPubSub{
-                        c:    c,
-                        sub:  gcps.Subscription("testsub"),
-                        msgC: make(chan *pubsub.Message),
+                src := &testSource{
+                        eventC: make(chan event),
                 }
+                ack := &testACK{}
                 s := Subscriber{
-                        PubSub:       ps,
-                        Subscription: ps.sub,
+                        S:       src,
+                        A:       ack,
+                        Workers: 8,
                 }
 
-                var received []string
-                var receivedMu sync.Mutex
-                cb := func(msg *pubsub.Message) {
-                        receivedMu.Lock()
-                        defer receivedMu.Unlock()
-                        received = append(received, msg.ID)
-                }
-
-                // If a subscriber goroutine is sleeping, advance to wake it.
-                tc.SetTimerCallback(func(d time.Duration, t clock.Timer) {
-                        tc.Add(d)
-                })
-
-                Convey(`A Subscriber can pull messages.`, func() {
+                var seenMu sync.Mutex
+                seen := map[string]struct{}{}
+                blacklist := map[string]struct{}{}
+                runWith := func(cb func()) {
                         doneC := make(chan struct{})
                         go func() {
                                 defer close(doneC)
-                                s.Run(c, cb)
+                                s.Run(c, func(msg *pubsub.Message) bool {
+                                        seenMu.Lock()
+                                        defer seenMu.Unlock()
+                                        seen[msg.ID] = struct{}{}
+
+                                        _, ok := blacklist[msg.ID]
+                                        return !ok
+                                })
                         }()
 
-                        var msgs []string
-                        for i := 0; i < 1024; i++ {
-                                msgs = append(msgs, fmt.Sprintf("%08x", i))
-                        }
-                        ps.send(msgs...)
-
+                        cb()
                         cancelFunc()
                         <-doneC
+                }
 
-                        sort.Strings(msgs)
-                        sort.Strings(received)
-                        So(received, ShouldResemble, msgs)
+                Convey(`A Subscriber can pull and ACK messages.`, func() {
+                        var msgs []string
+                        runWith(func() {
+                                for i := 0; i < 1024; i++ {
+                                        v := fmt.Sprintf("%08x", i)
+                                        msgs = append(msgs, v)
+                                        src.message(v)
+                                }
+                        })
+
+                        So(dumpStringSet(seen), ShouldResembleV, msgs)
+                        So(ack.getACKs(), ShouldResembleV, msgs)
                 })
 
-                Convey(`A Subscriber will continue retrying if there is a Pub/Sub error.`, func() {
-                        var errMu sync.Mutex
-                        count := 0
-                        ps.errCB = func() error {
-                                errMu.Lock()
-                                defer errMu.Unlock()
+                Convey(`A Subscriber that encounters an empty message set will sleep and try again.`, func() {
+                        var count int32
+                        tc.SetTimerCallback(func(d time.Duration, t clock.Timer) {
+                                if atomic.AddInt32(&count, 1) > 1 {
+                                        panic("should not have this many callbacks")
+                                }
+                                tc.Add(d)
+                        })
 
-                                if count < 1024 {
-                                        count++
-                                        return errors.WrapTransient(errors.New("test error"))
+                        runWith(func() {
+                                src.message("a", "b", "", "c", "d")
+                        })
+
+                        So(dumpStringSet(seen), ShouldResembleV, []string{"a", "b", "c", "d"})
+                        So(ack.getACKs(), ShouldResembleV, []string{"a", "b", "c", "d"})
+                })
+
+                Convey(`A Subscriber that encounters a Source error will sleep and try again.`, func() {
+                        var count int32
+                        tc.SetTimerCallback(func(d time.Duration, t clock.Timer) {
+                                if atomic.AddInt32(&count, 1) > 1 {
+                                        panic("should not have this many callbacks")
                                 }
-                                return nil
-                        }
+                                tc.Add(d)
+                        })
 
-                        doneC := make(chan struct{})
-                        go func() {
-                                defer close(doneC)
-                                s.Run(c, cb)
-                        }()
+                        runWith(func() {
+                                src.message("a", "b")
+                                src.error(errors.New("test error"))
+                                src.message("c", "d")
+                        })
 
-                        var msgs []string
-                        for i := 0; i < 1024; i++ {
-                                msgs = append(msgs, fmt.Sprintf("%08x", i))
-                        }
-                        ps.send(msgs...)
-
-                        cancelFunc()
-                        <-doneC
-
-                        sort.Strings(msgs)
-                        sort.Strings(received)
-                        So(received, ShouldResemble, msgs)
+                        So(dumpStringSet(seen), ShouldResembleV, []string{"a", "b", "c", "d"})
+                        So(ack.getACKs(), ShouldResembleV, []string{"a", "b", "c", "d"})
                 })
         })
 }