Repository: numaproj/numaflow
Pull request #661 (2023-04-03 15:41:11)

chore: generator : max inflight go routines should be 1
while testing back pressure scenarios, we need to make sure that the source handles back pressure appropriately.  while handling back pressure, we also need to make sure that all the records are generated and nothing is skipped. 



Patch:
diff --git a/pkg/sources/generator/tickgen.go b/pkg/sources/generator/tickgen.go
index 8554a7b3..4c6cabd6 100644
--- a/pkg/sources/generator/tickgen.go
+++ b/pkg/sources/generator/tickgen.go
@@ -24,7 +24,6 @@ import (
 	"encoding/json"
 	"fmt"
 	"strconv"
-	"sync/atomic"
 	"time"
 
 	"go.uber.org/zap"
@@ -312,12 +311,64 @@ func (mg *memgen) Start() <-chan struct{} {
 	return mg.forwarder.Start()
 }
 
+func (mg *memgen) NewWorker(ctx context.Context, rate int) func(chan time.Time, chan struct{}) {
+
+	return func(tickChan chan time.Time, done chan struct{}) {
+		for {
+			select {
+			case <-ctx.Done():
+				log.Info("Context.Done is called. emptying any pending ticks")
+				for len(tickChan) > 0 {
+					<-tickChan
+				}
+				close(done)
+				close(mg.srcchan)
+				return
+			case ts := <-tickChan:
+				tickgenSourceCount.With(map[string]string{metrics.LabelVertex: mg.name, metrics.LabelPipeline: mg.pipelineName})
+				// we would generate all the keys in a round robin fashion
+				// even if there are multiple pods, all the pods will generate same keys in the same order.
+				// TODO: alternatively, we could also think about generating a subset of keys per pod.
+				t := ts.UnixNano()
+				for i := 0; i < rate; i++ {
+					for k := int32(0); k < mg.keyCount; k++ {
+						key := fmt.Sprintf("key-%d", k)
+						payload := mg.genfn(mg.msgSize, mg.value, t)
+						r := record{data: payload, offset: time.Now().UTC().UnixNano(), key: key}
+						select {
+						case <-ctx.Done():
+							log.Info("Context.Done is called. returning from the inner function")
+							return
+						case mg.srcchan <- r:
+						}
+					}
+				}
+			}
+		}
+	}
+}
+
 // generator fires once per time unit and generates records and writes them to the channel
 func (mg *memgen) generator(ctx context.Context, rate int, timeunit time.Duration) {
 	go func() {
-		var rCount int32 = 0
-		// we are capping the limit at 10000 msgs / second
-		var limit = int32(10000 / rate)
+		// capping the rate to 10000 msgs/sec
+		if rate > 10000 {
+			log.Infow("Capping the rate to 10000 msg/sec. rate has been changed from %d to 10000", rate)
+			rate = 10000
+		}
+
+		tickChan := make(chan time.Time, 1000)
+		doneChan := make(chan struct{})
+		childCtx, childCancel := context.WithCancel(ctx)
+
+		defer childCancel()
+
+		// make sure that there is only one worker all the time.
+		// even when there is back pressure, max number of go routines inflight should be 1.
+		// at the same time, we dont want to miss any ticks that cannot be processed.
+		worker := mg.NewWorker(childCtx, rate)
+		go worker(tickChan, doneChan)
+
 		ticker := time.NewTicker(timeunit)
 		defer ticker.Stop()
 		for {
@@ -326,32 +377,11 @@ func (mg *memgen) generator(ctx context.Context, rate int, timeunit time.Duratio
 			// when context closes
 			case <-ctx.Done():
 				log.Info("Context.Done is called. exiting generator loop.")
+				childCancel()
+				<-doneChan
 				return
 			case ts := <-ticker.C:
-				tickgenSourceCount.With(map[string]string{metrics.LabelVertex: mg.name, metrics.LabelPipeline: mg.pipelineName})
-				// swapped implies that the rCount is at limit
-				if !atomic.CompareAndSwapInt32(&rCount, limit-1, limit) {
-					go func(t int64) {
-						atomic.AddInt32(&rCount, 1)
-						defer atomic.AddInt32(&rCount, -1)
-						// we would generate all the keys in a round robin fashion
-						// even if there are multiple pods, all the pods will generate same keys in the same order.
-						// alternatively, we could also think about generating a subset of keys per pod.
-						for i := 0; i < rate; i++ {
-							for k := int32(0); k < mg.keyCount; k++ {
-								key := fmt.Sprintf("key-%d", k)
-								payload := mg.genfn(mg.msgSize, mg.value, t)
-								r := record{data: payload, offset: time.Now().UTC().UnixNano(), key: key}
-								select {
-								case <-ctx.Done():
-									log.Info("Context.Done is called. returning from the inner function")
-									return
-								case mg.srcchan <- r:
-								}
-							}
-						}
-					}(ts.UnixNano())
-				}
+				tickChan <- ts
 			}
 		}
 	}()


Test patch:
diff --git a/pkg/sources/generator/tickgen_test.go b/pkg/sources/generator/tickgen_test.go
index 3465a5a8..9aa8c489 100644
--- a/pkg/sources/generator/tickgen_test.go
+++ b/pkg/sources/generator/tickgen_test.go
@@ -35,7 +35,7 @@ import (
 )
 
 func TestRead(t *testing.T) {
-	dest := simplebuffer.NewInMemoryBuffer("writer", 20, simplebuffer.WithReadTimeOut(10*time.Second))
+	dest := simplebuffer.NewInMemoryBuffer("writer", 100, simplebuffer.WithReadTimeOut(10*time.Second))
 	ctx := context.Background()
 	vertex := &dfv1.Vertex{
 		ObjectMeta: v1.ObjectMeta{
@@ -66,11 +66,21 @@ func TestRead(t *testing.T) {
 	assert.NoError(t, err)
 	_ = mgen.Start()
 
-	time.Sleep(time.Millisecond)
 	msgs, err := dest.Read(ctx, 5)
 
+	mgen.Stop()
+
 	assert.Nil(t, err)
 	assert.Equal(t, 5, len(msgs))
+
+	// wait for the context to be completely stopped.
+	for {
+		_, ok := <-mgen.srcchan
+		if !ok {
+			break
+		}
+	}
+
 }
 
 // Intention of this test is test the wiring for stop.
@@ -81,7 +91,8 @@ func TestStop(t *testing.T) {
 	// for use by the buffer reader on the other side of the stream
 	ctx := context.Background()
 
-	dest := simplebuffer.NewInMemoryBuffer("writer", 100)
+	// default rpu is 5. set the test to run for 2 ticks.
+	dest := simplebuffer.NewInMemoryBuffer("writer", 10)
 	vertex := &dfv1.Vertex{
 		ObjectMeta: v1.ObjectMeta{
 			Name: "memgen",
@@ -113,7 +124,7 @@ func TestStop(t *testing.T) {
 	starttime := time.Now()
 	// wait for some messages
 	for {
-		if dest.IsFull() {
+		if !dest.IsEmpty() {
 			break
 		}
 	}
@@ -121,11 +132,8 @@ func TestStop(t *testing.T) {
 
 	t.Logf("took [%v] millis to detect IsFull ", duration.Milliseconds())
 
-	rctx, rcf := context.WithCancel(ctx)
-
 	// initiate shutdown
 	mgen.Stop()
-	rcf()
 
 	// reader should have drained all the messages
 	// try to read everything
@@ -140,10 +148,10 @@ func TestStop(t *testing.T) {
 			assert.Greater(t, msgsread, 0)
 			return
 		default:
-			msgs, rerr := dest.Read(rctx, 1)
+			msgs, rerr := dest.Read(ctx, 1)
 			assert.Nil(t, rerr)
 			if len(msgs) > 0 {
-				_ = dest.Ack(rctx, []isb.Offset{msgs[0].ReadOffset})
+				_ = dest.Ack(ctx, []isb.Offset{msgs[0].ReadOffset})
 			}
 		}
 		msgsread += 1