Merge branch 'master' into peter/upgrade-lz4

Author: j1010001

engine/notifier_test.go

@@ -1,12 +1,13 @@
 package engine
 
 import (
+	"context"
 	"fmt"
 	"sync"
 	"testing"
-	"time"
+	"testing/synctest"
 
-	"github.com/stretchr/testify/require"
+	"github.com/stretchr/testify/assert"
 	"go.uber.org/atomic"
 )
 
@@ -17,8 +18,8 @@ func TestNotifier_PassByValue(t *testing.T) {
 
 	var sent sync.WaitGroup
 	sent.Add(1)
-	go func(n Notifier) {
-		notifier.Notify()
+	go func(passedByValue Notifier) {
+		passedByValue.Notify()
 		sent.Done()
 	}(notifier)
 	sent.Wait()
@@ -38,7 +39,7 @@ func TestNotifier_NoNotificationsInitialization(t *testing.T) {
 	select {
 	case <-notifier.Channel():
 		t.Fail()
-	default: //expected
+	default: // expected
 	}
 }
 
@@ -52,12 +53,10 @@ func TestNotifier_ManyNotifications(t *testing.T) {
 	notifier := NewNotifier()
 
 	var counter sync.WaitGroup
-	for i := 0; i < 10; i++ {
-		counter.Add(1)
-		go func() {
+	for range 10 {
+		counter.Go(func() {
 			notifier.Notify()
-			counter.Done()
-		}()
+		})
 	}
 	counter.Wait()
 
@@ -67,141 +66,132 @@ func TestNotifier_ManyNotifications(t *testing.T) {
 	select {
 	case <-c: // expected
 	default:
-		t.Fail()
+		t.Error("expected one notification to be available")
 	}
 
-	// attempt to consume first notification
+	// attempt to consume second notification
 	// expect that no notification is available
 	select {
 	case <-c:
-		t.Fail()
-	default: //expected
+		t.Error("expected only one notification to be available")
+	default: // expected
 	}
 }
 
-// TestNotifier_ManyConsumers spans many worker routines and
-// sends just as many notifications with small delays. We require that
-// all workers eventually get a notification.
+// TestNotifier_ManyConsumers spans many worker routines and sends just as many notifications.
+// We require that all workers eventually get a notification.
 func TestNotifier_ManyConsumers(t *testing.T) {
 	singleTestRun := func(t *testing.T) {
 		t.Parallel()
-		notifier := NewNotifier()
-		c := notifier.Channel()
-
-		// spawn 100 worker routines to each wait for a notification
-		var startingWorkers sync.WaitGroup
-		pendingWorkers := atomic.NewInt32(100)
-		for i := 0; i < 100; i++ {
-			startingWorkers.Add(1)
-			go func() {
-				startingWorkers.Done()
-				<-c
-				pendingWorkers.Dec()
-			}()
-		}
-		startingWorkers.Wait()
-
-		// send 100 notifications, with small delays
-		for i := 0; i < 100; i++ {
-			notifier.Notify()
-			time.Sleep(100 * time.Millisecond)
-		}
+		synctest.Test(t, func(t *testing.T) {
+			notifier := NewNotifier()
+			c := notifier.Channel()
+
+			// spawn 100 worker routines to each wait for a notification
+			pendingWorkers := atomic.NewInt32(100)
+			for range 100 {
+				go func() {
+					<-c
+					pendingWorkers.Dec()
+				}()
+			}
+
+			// wait until all workers are blocked on the notifier channel
+			synctest.Wait()
+
+			for range 100 {
+				notifier.Notify()
+
+				// wait for the previous notification to be consumed to ensure that the producer
+				// won't drop any notifications.
+				// NOTE: this is necessary because golang channels do not provide atomic consistency.
+				// Specifically, it means that when we send a notification while workers are waiting
+				// it is not guaranteed that a worker will atomically receive that notification. In
+				// other words, the channel might behave as if there was no worker waiting and de-
+				// duplicating notifications. For details, see
+				// https://www.notion.so/flowfoundation/Golang-Channel-Consistency-19a1aee12324817699b1ff162921d8fc
+				synctest.Wait()
+			}
+
+			// wait until all workers are done
+			synctest.Wait()
 
-		// require that all workers got a notification
-		if !conditionEventuallySatisfied(func() bool { return pendingWorkers.Load() == 0 }, 3*time.Second, 100*time.Millisecond) {
-			require.Fail(t, "timed out", "still awaiting %d workers to get notification", pendingWorkers.Load())
-		}
+			// require that all workers got a notification
+			assert.Equal(t, int32(0), pendingWorkers.Load())
+		})
 	}
 
-	for r := 0; r < 100; r++ {
+	for r := range 100 {
 		t.Run(fmt.Sprintf("run %d", r), singleTestRun)
 	}
 }
 
 // TestNotifier_AllWorkProcessed spans many routines pushing work and fewer
-// routines consuming work. We require that all worker is eventually processed.
+// routines consuming work. We require that all work is eventually processed.
 func TestNotifier_AllWorkProcessed(t *testing.T) {
 	singleTestRun := func(t *testing.T) {
 		t.Parallel()
-		notifier := NewNotifier()
-
-		totalWork := int32(100)
-		pendingWorkQueue := make(chan struct{}, totalWork)
-		scheduledWork := atomic.NewInt32(0)
-		consumedWork := atomic.NewInt32(0)
-
-		// starts the consuming first, because if we starts the production first instead, then
-		// we might finish pushing all jobs, before any of our consumer has started listening
-		// to the queue.
-		var consumersAllReady sync.WaitGroup
-		consumersAllReady.Add(5)
-
-		// 5 routines consuming work
-		for i := 0; i < 5; i++ {
-			go func() {
-				consumersAllReady.Done()
-				for consumedWork.Load() < totalWork {
-					<-notifier.Channel()
-				L:
+		synctest.Test(t, func(t *testing.T) {
+			ctx, cancel := context.WithCancel(context.Background())
+
+			notifier := NewNotifier()
+
+			producerCount := int32(10) // number of producers
+			producerJobs := int32(10)  // number of tasks that each producer will queue up
+			pendingWorkQueue := make(chan struct{}, producerCount*producerJobs)
+			consumedWork := atomic.NewInt32(0)
+
+			// start the consumers first, otherwise we might finish pushing all jobs, before any of
+			// our consumer has started listening to the queue.
+
+			processAllPending := func() {
+				for {
+					select {
+					case <-pendingWorkQueue:
+						consumedWork.Inc()
+					default:
+						return
+					}
+				}
+			}
+
+			// 5 routines consuming work
+			for range 5 {
+				go func() {
 					for {
 						select {
-						case <-pendingWorkQueue:
-							consumedWork.Inc()
-						default:
-							break L
+						case <-ctx.Done():
+							return
+						case <-notifier.Channel():
+							processAllPending()
 						}
 					}
-				}
-			}()
-		}
-
-		// wait long enough for all consumer to be ready for new notification.
-		consumersAllReady.Wait()
-
-		var workersAllReady sync.WaitGroup
-		workersAllReady.Add(10)
-
-		// 10 routines pushing work
-		for i := 0; i < 10; i++ {
-			go func() {
-				workersAllReady.Done()
-				for scheduledWork.Inc() <= totalWork {
-					pendingWorkQueue <- struct{}{}
-					notifier.Notify()
-				}
-			}()
-		}
-
-		// wait long enough for all workers to be started.
-		workersAllReady.Wait()
+				}()
+			}
 
-		// require that all work is eventually consumed
-		if !conditionEventuallySatisfied(func() bool { return consumedWork.Load() == totalWork }, 3*time.Second, 100*time.Millisecond) {
-			require.Fail(t, "timed out", "only consumed %d units of work but expecting %d", consumedWork.Load(), totalWork)
-		}
-	}
+			// wait for all consumers to block on the notifier channel
+			synctest.Wait()
 
-	for r := 0; r < 100; r++ {
-		t.Run(fmt.Sprintf("run %d", r), singleTestRun)
-	}
-}
+			for range producerCount {
+				go func() {
+					for range producerJobs {
+						pendingWorkQueue <- struct{}{}
+						notifier.Notify()
+					}
+				}()
+			}
 
-func conditionEventuallySatisfied(condition func() bool, waitFor time.Duration, tick time.Duration) bool {
-	done := make(chan struct{})
+			// wait for all producers and consumers to block. at this point, all jobs should be consumed.
+			synctest.Wait()
+			assert.Equal(t, producerCount*producerJobs, consumedWork.Load())
 
-	go func() {
-		for range time.Tick(tick) {
-			if condition() {
-				close(done)
-				return
-			}
-		}
-	}()
+			// shutdown blocked consumers and wait for them to complete
+			cancel()
+			synctest.Wait()
+		})
+	}
 
-	select {
-	case <-time.After(waitFor):
-		return false
-	case <-done:
-		return true
+	for r := range 100 {
+		t.Run(fmt.Sprintf("run %d", r), singleTestRun)
 	}
 }