Repository: numaproj/numaflow
Pull request #419 (2022-12-09 23:23:32)

chore: refactor reduce ctx
Signed-off-by: Vigith Maurice <vigith@gmail.com>

The context passed to `forwardAChunk` will never be canceled in an error loop scenario 


Patch:
diff --git a/pkg/reduce/readloop/ordered.go b/pkg/reduce/readloop/ordered.go
index 6ae7757f..a7209105 100644
--- a/pkg/reduce/readloop/ordered.go
+++ b/pkg/reduce/readloop/ordered.go
@@ -55,16 +55,15 @@ type orderedForwarder struct {
 
 // newOrderedForwarder returns an orderedForwarder.
 func newOrderedForwarder(ctx context.Context) *orderedForwarder {
-	return &orderedForwarder{
+	of := &orderedForwarder{
 		taskDone:  make(chan struct{}),
 		taskQueue: list.New(),
 		log:       logging.FromContext(ctx),
 	}
-}
 
-// startUp starts forwarder.
-func (of *orderedForwarder) startUp(ctx context.Context) {
 	go of.forward(ctx)
+
+	return of
 }
 
 // schedulePnF creates and schedules the PnF routine.
@@ -109,9 +108,10 @@ func (of *orderedForwarder) reduceOp(ctx context.Context, t *task) {
 		of.log.Errorw("Process failed", zap.String("partitionID", t.pf.PartitionID.String()), zap.Error(err))
 		time.Sleep(retryDelay)
 	}
-	// after retrying indicate that we are done with processing the package. the processing can move on
+	// indicate that we are done with reduce UDF invocation.
 	close(t.doneCh)
 	of.log.Debugw("Process->Reduce call took ", zap.String("partitionID", t.pf.PartitionID.String()), zap.Int64("duration(ms)", time.Since(start).Milliseconds()))
+
 	// notify that some work has been completed
 	select {
 	case of.taskDone <- struct{}{}:
diff --git a/pkg/reduce/readloop/readloop.go b/pkg/reduce/readloop/readloop.go
index cf77d5a8..442e1b8e 100644
--- a/pkg/reduce/readloop/readloop.go
+++ b/pkg/reduce/readloop/readloop.go
@@ -67,8 +67,7 @@ func NewReadLoop(ctx context.Context,
 	windowingStrategy window.Windower,
 	toBuffers map[string]isb.BufferWriter,
 	whereToDecider forward.ToWhichStepDecider,
-	pw map[string]publish.Publisher,
-) *ReadLoop {
+	pw map[string]publish.Publisher) (*ReadLoop, error) {
 
 	op := newOrderedForwarder(ctx)
 
@@ -82,8 +81,10 @@ func NewReadLoop(ctx context.Context,
 		whereToDecider:   whereToDecider,
 		publishWatermark: pw,
 	}
-	op.startUp(ctx)
-	return rl
+
+	err := rl.Startup(ctx)
+
+	return rl, err
 }
 
 // Startup starts up the read-loop, because during boot up, it has to replay the data from the persistent store of
diff --git a/pkg/reduce/reduce.go b/pkg/reduce/reduce.go
index eebe50f2..1a925977 100644
--- a/pkg/reduce/reduce.go
+++ b/pkg/reduce/reduce.go
@@ -39,6 +39,7 @@ import (
 
 // DataForward reads data from isb and forwards them to readloop
 type DataForward struct {
+	ctx                 context.Context
 	vertexInstance      *dfv1.VertexInstance
 	fromBuffer          isb.BufferReader
 	toBuffers           map[string]isb.BufferWriter
@@ -70,8 +71,9 @@ func NewDataForward(ctx context.Context,
 		}
 	}
 
-	rl := readloop.NewReadLoop(ctx, udf, pbqManager, windowingStrategy, toBuffers, whereToDecider, watermarkPublishers)
-	return &DataForward{
+	rl, err := readloop.NewReadLoop(ctx, udf, pbqManager, windowingStrategy, toBuffers, whereToDecider, watermarkPublishers)
+	df := &DataForward{
+		ctx:                 ctx,
 		vertexInstance:      vertexInstance,
 		fromBuffer:          fromBuffer,
 		toBuffers:           toBuffers,
@@ -80,22 +82,17 @@ func NewDataForward(ctx context.Context,
 		watermarkPublishers: watermarkPublishers,
 		windowingStrategy:   windowingStrategy,
 		log:                 logging.FromContext(ctx),
-		opts:                options}, nil
+		opts:                options}
+
+	return df, err
 }
 
 // Start starts forwarding messages to readloop
-func (d *DataForward) Start(ctx context.Context) {
-	err := d.readloop.Startup(ctx)
-	if err != nil {
-		d.log.Errorw("Failed to start the data forwarder in reduce vertex", zap.Error(err))
-	}
-	rCtx, rCancel := context.WithCancel(context.Background())
+func (d *DataForward) Start() {
 	for {
 		select {
-		case <-ctx.Done():
+		case <-d.ctx.Done():
 			d.log.Infow("Stopping reduce data forwarder...")
-			cctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
-			defer cancel()
 
 			if err := d.fromBuffer.Close(); err != nil {
 				d.log.Errorw("Failed to close buffer reader, shutdown anyways...", zap.Error(err))
@@ -103,8 +100,9 @@ func (d *DataForward) Start(ctx context.Context) {
 				d.log.Infow("Closed buffer reader", zap.String("bufferFrom", d.fromBuffer.GetName()))
 			}
 
-			// cancel the readloop context so that all the messages are persisted and acked.
-			rCancel()
+			// hard shutdown after timeout
+			cctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+			defer cancel()
 			// allow readloop to clean itself up.
 			d.readloop.ShutDown(cctx)
 
@@ -132,7 +130,7 @@ func (d *DataForward) Start(ctx context.Context) {
 			// pass the child context so that the reader can be closed before the readloop
 			// this way we can avoid the race condition and have all the read messages persisted
 			// and acked.
-			d.forwardAChunk(rCtx)
+			d.forwardAChunk(d.ctx)
 		}
 	}
 }
diff --git a/pkg/udf/reduce_udf.go b/pkg/udf/reduce_udf.go
index 7951b158..6af0b63e 100644
--- a/pkg/udf/reduce_udf.go
+++ b/pkg/udf/reduce_udf.go
@@ -175,7 +175,7 @@ func (u *ReduceUDFProcessor) Start(ctx context.Context) error {
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
-		dataForwarder.Start(ctx)
+		dataForwarder.Start()
 		log.Info("Forwarder stopped, exiting reduce udf data processor...")
 	}()
 


Test patch:
diff --git a/pkg/reduce/readloop/ordered_test.go b/pkg/reduce/readloop/ordered_test.go
index 63c0f1a7..40c8e417 100644
--- a/pkg/reduce/readloop/ordered_test.go
+++ b/pkg/reduce/readloop/ordered_test.go
@@ -26,13 +26,14 @@ import (
 	"github.com/numaproj/numaflow/pkg/isb/testutils"
 	"github.com/numaproj/numaflow/pkg/watermark/generic"
 
+	"github.com/stretchr/testify/assert"
+
 	dfv1 "github.com/numaproj/numaflow/pkg/apis/numaflow/v1alpha1"
 	"github.com/numaproj/numaflow/pkg/isb"
 	"github.com/numaproj/numaflow/pkg/pbq"
 	"github.com/numaproj/numaflow/pkg/pbq/partition"
 	"github.com/numaproj/numaflow/pkg/pbq/store/memory"
 	"github.com/numaproj/numaflow/pkg/udf/applier"
-	"github.com/stretchr/testify/assert"
 )
 
 type myForwardTest struct {
@@ -104,7 +105,6 @@ func TestOrderedProcessing(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			// clean out the task queue before we start a run
 			op := newOrderedForwarder(ctx)
-			op.startUp(ctx)
 			// although this could be declared outside, since we are using common naming scheme for partitions,
 			// things will go haywire.
 			pbqManager, _ := pbq.NewManager(ctx, memory.NewMemoryStores(memory.WithStoreSize(100)),
diff --git a/pkg/reduce/readloop/readloop_test.go b/pkg/reduce/readloop/readloop_test.go
index 9f453c4c..34403d50 100644
--- a/pkg/reduce/readloop/readloop_test.go
+++ b/pkg/reduce/readloop/readloop_test.go
@@ -7,6 +7,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
+
 	"github.com/numaproj/numaflow/pkg/isb"
 	"github.com/numaproj/numaflow/pkg/isb/stores/simplebuffer"
 	"github.com/numaproj/numaflow/pkg/pbq"
@@ -14,7 +16,6 @@ import (
 	"github.com/numaproj/numaflow/pkg/pbq/store/memory"
 	"github.com/numaproj/numaflow/pkg/watermark/generic"
 	"github.com/numaproj/numaflow/pkg/window/strategy/fixed"
-	"github.com/stretchr/testify/assert"
 )
 
 // PayloadForTest is a dummy payload for testing.
@@ -26,11 +27,11 @@ type PayloadForTest struct {
 type SumReduceTest struct {
 }
 
-func (s *SumReduceTest) WhereTo(s2 string) ([]string, error) {
+func (s *SumReduceTest) WhereTo(_ string) ([]string, error) {
 	return []string{"reduce-buffer"}, nil
 }
 
-func (s *SumReduceTest) ApplyReduce(ctx context.Context, partitionID *partition.ID, messageStream <-chan *isb.ReadMessage) ([]*isb.Message, error) {
+func (s *SumReduceTest) ApplyReduce(_ context.Context, partitionID *partition.ID, messageStream <-chan *isb.ReadMessage) ([]*isb.Message, error) {
 	sum := 0
 	for msg := range messageStream {
 		var payload PayloadForTest
@@ -118,7 +119,7 @@ func TestReadLoop_Startup(t *testing.T) {
 
 	window := fixed.NewFixed(60 * time.Second)
 
-	rl := NewReadLoop(ctx, &SumReduceTest{}, pManager, window, toSteps, &SumReduceTest{}, pw)
+	rl, _ := NewReadLoop(ctx, &SumReduceTest{}, pManager, window, toSteps, &SumReduceTest{}, pw)
 
 	err := rl.Startup(ctx)
 	assert.NoError(t, err)
@@ -179,7 +180,7 @@ func TestReadLoop_Startup(t *testing.T) {
 
 }
 
-func createStoreMessages(ctx context.Context, key string, value int, eventTime time.Time, count int) []*isb.ReadMessage {
+func createStoreMessages(_ context.Context, key string, value int, eventTime time.Time, count int) []*isb.ReadMessage {
 	readMessages := make([]*isb.ReadMessage, count)
 	for j := 0; j < count; j++ {
 		result, _ := json.Marshal(PayloadForTest{
diff --git a/pkg/reduce/reduce_test.go b/pkg/reduce/reduce_test.go
index 2085228d..b88d5d59 100644
--- a/pkg/reduce/reduce_test.go
+++ b/pkg/reduce/reduce_test.go
@@ -105,7 +105,7 @@ type CounterReduceTest struct {
 }
 
 // Reduce returns a result with the count of messages
-func (f CounterReduceTest) ApplyReduce(ctx context.Context, partitionID *partition.ID, messageStream <-chan *isb.ReadMessage) ([]*isb.Message, error) {
+func (f CounterReduceTest) ApplyReduce(_ context.Context, partitionID *partition.ID, messageStream <-chan *isb.ReadMessage) ([]*isb.Message, error) {
 	count := 0
 	for range messageStream {
 		count += 1
@@ -130,14 +130,14 @@ func (f CounterReduceTest) ApplyReduce(ctx context.Context, partitionID *partiti
 	}, nil
 }
 
-func (f CounterReduceTest) WhereTo(s string) ([]string, error) {
+func (f CounterReduceTest) WhereTo(_ string) ([]string, error) {
 	return []string{"reduce-to-buffer"}, nil
 }
 
 type SumReduceTest struct {
 }
 
-func (s SumReduceTest) ApplyReduce(ctx context.Context, partitionID *partition.ID, messageStream <-chan *isb.ReadMessage) ([]*isb.Message, error) {
+func (s SumReduceTest) ApplyReduce(_ context.Context, partitionID *partition.ID, messageStream <-chan *isb.ReadMessage) ([]*isb.Message, error) {
 	sum := 0
 	for msg := range messageStream {
 		var payload PayloadForTest
@@ -167,7 +167,7 @@ func (s SumReduceTest) ApplyReduce(ctx context.Context, partitionID *partition.I
 type MaxReduceTest struct {
 }
 
-func (m MaxReduceTest) ApplyReduce(ctx context.Context, partitionID *partition.ID, messageStream <-chan *isb.ReadMessage) ([]*isb.Message, error) {
+func (m MaxReduceTest) ApplyReduce(_ context.Context, partitionID *partition.ID, messageStream <-chan *isb.ReadMessage) ([]*isb.Message, error) {
 	mx := math.MinInt64
 	for msg := range messageStream {
 		var payload PayloadForTest
@@ -244,7 +244,7 @@ func TestDataForward_StartWithNoOpWM(t *testing.T) {
 		window, WithReadBatchSize(10))
 	assert.NoError(t, err)
 
-	go reduceDataForwarder.Start(child)
+	go reduceDataForwarder.Start()
 
 	for to.IsEmpty() {
 		select {
@@ -314,7 +314,7 @@ func TestReduceDataForward_Count(t *testing.T) {
 	assert.NoError(t, err)
 
 	// start the forwarder
-	go reduceDataForward.Start(ctx)
+	go reduceDataForward.Start()
 
 	// start the producer
 	go publishMessages(ctx, startTime, messageValue, 300, 10, p[fromBuffer.GetName()], fromBuffer)
@@ -387,7 +387,7 @@ func TestReduceDataForward_Sum(t *testing.T) {
 	assert.NoError(t, err)
 
 	// start the forwarder
-	go reduceDataForward.Start(ctx)
+	go reduceDataForward.Start()
 
 	// start the producer
 	go publishMessages(ctx, startTime, messageValue, 300, 10, p[fromBuffer.GetName()], fromBuffer)
@@ -460,7 +460,7 @@ func TestReduceDataForward_Max(t *testing.T) {
 	assert.NoError(t, err)
 
 	// start the forwarder
-	go reduceDataForward.Start(ctx)
+	go reduceDataForward.Start()
 
 	// start the producer
 	go publishMessages(ctx, startTime, messageValue, 600, 10, p[fromBuffer.GetName()], fromBuffer)
@@ -535,7 +535,7 @@ func TestReduceDataForward_SumWithDifferentKeys(t *testing.T) {
 	assert.NoError(t, err)
 
 	// start the forwarder
-	go reduceDataForward.Start(ctx)
+	go reduceDataForward.Start()
 
 	// start the producer
 	go publishMessages(ctx, startTime, messages, 600, 10, p[fromBuffer.GetName()], fromBuffer)
@@ -614,7 +614,7 @@ func TestReduceDataForward_NonKeyed(t *testing.T) {
 	assert.NoError(t, err)
 
 	// start the forwarder
-	go reduceDataForward.Start(ctx)
+	go reduceDataForward.Start()
 
 	// start the producer
 	go publishMessages(ctx, startTime, messages, 600, 10, p[fromBuffer.GetName()], fromBuffer)
@@ -694,7 +694,7 @@ func TestDataForward_WithContextClose(t *testing.T) {
 	assert.NoError(t, err)
 
 	// start the forwarder
-	go reduceDataForward.Start(cctx)
+	go reduceDataForward.Start()
 	// window duration is 300s, we are sending only 200 messages with event time less than window end time, so the window will not be closed
 	publishMessages(cctx, startTime, messages, 200, 10, p[fromBuffer.GetName()], fromBuffer)
 	// wait for the partitions to be created