| // Copyright 2019 Google LLC |
| // |
| // 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 scheduler_test |
| |
| import ( |
| "fmt" |
| "testing" |
| "time" |
| |
| "cloud.google.com/go/pubsub/internal/scheduler" |
| ) |
| |
| func TestPublishScheduler_Put_Basic(t *testing.T) { |
| done := make(chan struct{}) |
| defer close(done) |
| |
| keysHandled := map[string]chan int{} |
| handle := func(itemi interface{}) { |
| items := itemi.([]pair) |
| for _, item := range items { |
| keysHandled[item.k] <- item.v |
| } |
| } |
| s := scheduler.NewPublishScheduler(2, handle) |
| defer s.FlushAndStop() |
| |
| // If these values are too high, the race detector will fail with |
| // "race: limit on 8128 simultaneously alive goroutines is exceeded, dying". |
| numItems := 100 |
| numKeys := 10 |
| |
| for ki := 0; ki < numKeys; ki++ { |
| k := fmt.Sprintf("some_key_%d", ki) |
| keysHandled[k] = make(chan int, numItems) |
| } |
| |
| for ki := 0; ki < numKeys; ki++ { |
| k := fmt.Sprintf("some_key_%d", ki) |
| go func() { |
| for i := 0; i < numItems; i++ { |
| select { |
| case <-done: |
| return |
| default: |
| } |
| if err := s.Add(k, pair{k, i}, 1); err != nil { |
| t.Error(err) |
| } |
| } |
| }() |
| } |
| |
| for ki := 0; ki < numKeys; ki++ { |
| k := fmt.Sprintf("some_key_%d", ki) |
| for want := 0; want < numItems; want++ { |
| select { |
| case got := <-keysHandled[k]: |
| if got != want { |
| t.Fatalf("%s: got %d, want %d", k, got, want) |
| } |
| case <-time.After(5 * time.Second): |
| t.Fatalf("%s: expected key %s - item %d to be handled but never was", k, k, want) |
| } |
| } |
| } |
| } |
| |
| // Scheduler schedules many items of one key in order even when there are |
| // many workers. |
| func TestPublishScheduler_Put_ManyWithOneKey(t *testing.T) { |
| done := make(chan struct{}) |
| defer close(done) |
| |
| recvd := make(chan int) |
| handle := func(itemi interface{}) { |
| items := itemi.([]int) |
| for _, item := range items { |
| recvd <- item |
| } |
| } |
| s := scheduler.NewPublishScheduler(10, handle) |
| defer s.FlushAndStop() |
| |
| // If these values are too high, the race detector will fail with |
| // "race: limit on 8128 simultaneously alive goroutines is exceeded, dying". |
| numItems := 1000 |
| |
| go func() { |
| for i := 0; i < numItems; i++ { |
| select { |
| case <-done: |
| return |
| default: |
| } |
| if err := s.Add("some-key", i, 1); err != nil { |
| t.Error(err) |
| } |
| } |
| }() |
| |
| for want := 0; want < numItems; want++ { |
| select { |
| case got := <-recvd: |
| if got != want { |
| t.Fatalf("got %d, want %d", got, want) |
| } |
| case <-time.After(5 * time.Second): |
| t.Fatalf("timed out waiting for item %d to be handled", want) |
| } |
| } |
| } |
| |
| func TestPublishScheduler_DoesntRaceWithPublisher(t *testing.T) { |
| done := make(chan struct{}) |
| defer close(done) |
| |
| keysHandled := map[string]chan int{} |
| handle := func(itemi interface{}) { |
| items := itemi.([]pair) |
| for _, item := range items { |
| keysHandled[item.k] <- item.v |
| } |
| } |
| s := scheduler.NewPublishScheduler(2, handle) |
| defer s.FlushAndStop() |
| |
| // If these values are too high, the race detector will fail with |
| // "race: limit on 8128 simultaneously alive goroutines is exceeded, dying". |
| numItems := 100 |
| numKeys := 10 |
| |
| for ki := 0; ki < numKeys; ki++ { |
| k := fmt.Sprintf("some_key_%d", ki) |
| keysHandled[k] = make(chan int, numItems) |
| } |
| |
| for ki := 0; ki < numKeys; ki++ { |
| k := fmt.Sprintf("some_key_%d", ki) |
| go func() { |
| for i := 0; i < numItems; i++ { |
| select { |
| case <-done: |
| return |
| default: |
| } |
| if err := s.Add(k, pair{k, i}, 1); err != nil { |
| t.Error(err) |
| } |
| } |
| }() |
| } |
| |
| for ki := 0; ki < numKeys; ki++ { |
| k := fmt.Sprintf("some_key_%d", ki) |
| for want := 0; want < numItems; want++ { |
| select { |
| case got := <-keysHandled[k]: |
| if got != want { |
| t.Fatalf("%s: got %d, want %d", k, got, want) |
| } |
| case <-time.After(5 * time.Second): |
| t.Fatalf("%s: expected key %s - item %d to be handled but never was", k, k, want) |
| } |
| } |
| } |
| } |
| |
| // FlushAndStop blocks until all messages are processed. |
| func TestPublishScheduler_FlushAndStop(t *testing.T) { |
| for _, tc := range []struct { |
| name string |
| input map[string][]int |
| }{ |
| { |
| name: "two messages with the same key", |
| input: map[string][]int{"foo": {1, 2}}, |
| }, |
| { |
| name: "two messages with different keys", |
| input: map[string][]int{"foo": {1}, "bar": {2}}, |
| }, |
| { |
| name: "two messages with no key", |
| input: map[string][]int{"": {1, 2}}, |
| }, |
| } { |
| t.Run(tc.name, func(t *testing.T) { |
| recvd := make(chan int) |
| handle := func(itemi interface{}) { |
| for _, v := range itemi.([]int) { |
| recvd <- v |
| } |
| } |
| s := scheduler.NewPublishScheduler(1, handle) |
| for k, vs := range tc.input { |
| for _, v := range vs { |
| if err := s.Add(k, v, 1); err != nil { |
| t.Fatal(err) |
| } |
| } |
| } |
| |
| doneFlushing := make(chan struct{}) |
| go func() { |
| s.FlushAndStop() |
| close(doneFlushing) |
| }() |
| |
| time.Sleep(10 * time.Millisecond) |
| |
| select { |
| case <-doneFlushing: |
| t.Fatal("expected FlushAndStop to block until all messages handled, but it didn't") |
| default: |
| } |
| |
| select { |
| case <-recvd: |
| case <-time.After(time.Second): |
| t.Fatal("timed out waiting for first message to arrive") |
| } |
| |
| select { |
| case <-recvd: |
| case <-time.After(time.Second): |
| t.Fatal("timed out waiting for second message to arrive") |
| } |
| |
| select { |
| case <-doneFlushing: |
| case <-time.After(time.Second): |
| t.Fatal("timed out waiting for FlushAndStop to finish blocking") |
| } |
| }) |
| } |
| } |