LogDog: BigTable batching schema.

common/gcloud/pubsub/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"
@@ skipping 13 matching lines @@
 type event struct {
         msg *pubsub.Message
         err error
 }
 
 type testSource struct {
         sub    pubsub.Subscription
         eventC chan event
 }
 
-func (s *testSource) Pull(c context.Context) ([]*pubsub.Message, error) {
+func (s *testSource) Pull(c context.Context, batch int) ([]*pubsub.Message, error) {
+        var e event
+
         select {
         case <-c.Done():
-                return nil, c.Err()
+                // Enforce determinism, preferring events.
+                select {
+                case e = <-s.eventC:
+                        break
+                default:
+                        return nil, c.Err()
+                }
 
-        case e := <-s.eventC:
-                switch {
-                case e.err != nil:
-                        return nil, e.err
+        case e = <-s.eventC:
+                break
+        }
 
-                case e.msg != nil:
-                        return []*pubsub.Message{e.msg}, nil
+        switch {
+        case e.err != nil:
+                return nil, e.err
 
-                default:
-                        return nil, nil
-                }
+        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),
@@ skipping 47 matching lines @@
                 c, tc := testclock.UseTime(c, testclock.TestTimeLocal)
 
                 c, cancelFunc := context.WithCancel(c)
                 defer cancelFunc()
 
                 src := &testSource{
                         eventC: make(chan event),
                 }
                 ack := &testACK{}
                 s := Subscriber{
-                        S:           src,
-                        A:           ack,
-                        PullWorkers: 8,
+                        S:       src,
+                        A:       ack,
+                        Workers: 8,
                 }
 
+                // If not nil, processing goroutines will block on reads from this
+                // channel, one per message.
+                var pullC chan string
+
                 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, func(msg *pubsub.Message) bool {
+                                        if pullC != nil {
+                                                pullC <- msg.ID
+                                        }
+
                                         seenMu.Lock()
                                         defer seenMu.Unlock()
                                         seen[msg.ID] = struct{}{}
 
                                         _, ok := blacklist[msg.ID]
                                         return !ok
                                 })
                         }()
 
                         cb()
                         cancelFunc()
                         <-doneC
                 }
 
                 Convey(`A Subscriber can pull and ACK messages.`, func() {
                         var msgs []string
+                        pullC = make(chan string)
                         runWith(func() {
                                 for i := 0; i < 1024; i++ {
                                         v := fmt.Sprintf("%08x", i)
                                         msgs = append(msgs, v)
                                         src.message(v)
+
+                                        <-pullC
                                 }
                         })
 
                         So(dumpStringSet(seen), ShouldResemble, msgs)
                         So(ack.getACKs(), ShouldResemble, msgs)
                 })
 
                 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) {
@@ skipping 24 matching lines @@
                                 src.message("a", "b")
                                 src.error(errors.New("test error"))
                                 src.message("c", "d")
                         })
 
                         So(dumpStringSet(seen), ShouldResemble, []string{"a", "b", "c", "d"})
                         So(ack.getACKs(), ShouldResemble, []string{"a", "b", "c", "d"})
                 })
         })
 }