Repository: numaproj/numaflow
Pull request #433 (2022-12-19 07:28:22)

chore: optimizing ordered loop
optimization to the ordered loop so that queueing and dequeuing continues to stay at O(1). 

we can delay the enqueuing of the task until the close of the book as the desired order is guaranteed. 


Patch:
diff --git a/pkg/reduce/readloop/ordered.go b/pkg/reduce/readloop/ordered.go
index d608e668..afbf480e 100644
--- a/pkg/reduce/readloop/ordered.go
+++ b/pkg/reduce/readloop/ordered.go
@@ -22,9 +22,10 @@ import (
 	"sync"
 	"time"
 
-	"github.com/numaproj/numaflow/pkg/metrics"
 	"go.uber.org/zap"
 
+	"github.com/numaproj/numaflow/pkg/metrics"
+
 	"github.com/numaproj/numaflow/pkg/isb"
 	"github.com/numaproj/numaflow/pkg/isb/forward"
 	"github.com/numaproj/numaflow/pkg/watermark/publish"
@@ -71,34 +72,10 @@ func newOrderedForwarder(ctx context.Context, vertexName string, pipelineName st
 	return of
 }
 
-func (of *orderedForwarder) insertTask(tx context.Context, t *task) {
+func (of *orderedForwarder) insertTask(t *task) {
 	of.Lock()
 	defer of.Unlock()
-
-	if of.taskQueue.Len() == 0 {
-		of.taskQueue.PushBack(t)
-		return
-	}
-
-	earliestTask := of.taskQueue.Front().Value.(*task)
-	recentTask := of.taskQueue.Back().Value.(*task)
-
-	// late arrival
-	if !earliestTask.pf.PartitionID.End.Before(t.pf.PartitionID.End) {
-		of.taskQueue.PushFront(t)
-	} else if !recentTask.pf.PartitionID.End.After(t.pf.PartitionID.End) {
-		// early arrival
-		of.taskQueue.PushBack(t)
-	} else {
-		// a task in the middle
-		for e := of.taskQueue.Back(); e != nil; e = e.Prev() {
-			_task := e.Value.(*task)
-			if !_task.pf.PartitionID.End.After(t.pf.PartitionID.End) {
-				of.taskQueue.InsertAfter(t, e)
-				break
-			}
-		}
-	}
+	of.taskQueue.PushBack(t)
 }
 
 // schedulePnF creates and schedules the PnF routine.
@@ -108,7 +85,7 @@ func (of *orderedForwarder) schedulePnF(ctx context.Context,
 	partitionID partition.ID,
 	toBuffers map[string]isb.BufferWriter,
 	whereToDecider forward.ToWhichStepDecider,
-	pw map[string]publish.Publisher) {
+	pw map[string]publish.Publisher) *task {
 
 	pf := pnf.NewProcessAndForward(ctx, of.vertexName, of.pipelineName, partitionID, udf, pbq, toBuffers, whereToDecider, pw)
 	doneCh := make(chan struct{})
@@ -116,12 +93,10 @@ func (of *orderedForwarder) schedulePnF(ctx context.Context,
 		doneCh: doneCh,
 		pf:     pf,
 	}
-
-	of.insertTask(ctx, t)
 	partitionsInFlight.With(map[string]string{metrics.LabelVertex: of.vertexName, metrics.LabelPipeline: of.pipelineName}).Inc()
-
 	// invoke the reduce function
 	go of.reduceOp(ctx, t)
+	return t
 }
 
 // reduceOp invokes the reduce function. The reducer is a long running function since we stream in the data and it has
diff --git a/pkg/reduce/readloop/readloop.go b/pkg/reduce/readloop/readloop.go
index 520f8b65..49fb617f 100644
--- a/pkg/reduce/readloop/readloop.go
+++ b/pkg/reduce/readloop/readloop.go
@@ -33,10 +33,11 @@ import (
 	"sync"
 	"time"
 
-	"github.com/numaproj/numaflow/pkg/metrics"
 	"go.uber.org/zap"
 	"k8s.io/apimachinery/pkg/util/wait"
 
+	"github.com/numaproj/numaflow/pkg/metrics"
+
 	"github.com/numaproj/numaflow/pkg/isb"
 	"github.com/numaproj/numaflow/pkg/isb/forward"
 	"github.com/numaproj/numaflow/pkg/pbq"
@@ -51,16 +52,17 @@ import (
 
 // ReadLoop is responsible for reading and forwarding the message from ISB to PBQ.
 type ReadLoop struct {
-	vertexName       string
-	pipelineName     string
-	UDF              applier.ReduceApplier
-	pbqManager       *pbq.Manager
-	windower         window.Windower
-	op               *orderedForwarder
-	log              *zap.SugaredLogger
-	toBuffers        map[string]isb.BufferWriter
-	whereToDecider   forward.ToWhichStepDecider
-	publishWatermark map[string]publish.Publisher
+	vertexName            string
+	pipelineName          string
+	UDF                   applier.ReduceApplier
+	pbqManager            *pbq.Manager
+	windower              window.Windower
+	op                    *orderedForwarder
+	log                   *zap.SugaredLogger
+	toBuffers             map[string]isb.BufferWriter
+	whereToDecider        forward.ToWhichStepDecider
+	publishWatermark      map[string]publish.Publisher
+	udfInvocationTracking map[partition.ID]*task
 }
 
 // NewReadLoop initializes  and returns ReadLoop.
@@ -77,16 +79,17 @@ func NewReadLoop(ctx context.Context,
 	op := newOrderedForwarder(ctx, vertexName, pipelineName)
 
 	rl := &ReadLoop{
-		vertexName:       vertexName,
-		pipelineName:     pipelineName,
-		UDF:              udf,
-		windower:         windowingStrategy,
-		pbqManager:       pbqManager,
-		op:               op,
-		log:              logging.FromContext(ctx),
-		toBuffers:        toBuffers,
-		whereToDecider:   whereToDecider,
-		publishWatermark: pw,
+		vertexName:            vertexName,
+		pipelineName:          pipelineName,
+		UDF:                   udf,
+		windower:              windowingStrategy,
+		pbqManager:            pbqManager,
+		op:                    op,
+		log:                   logging.FromContext(ctx),
+		toBuffers:             toBuffers,
+		whereToDecider:        whereToDecider,
+		publishWatermark:      pw,
+		udfInvocationTracking: make(map[partition.ID]*task),
 	}
 
 	err := rl.Startup(ctx)
@@ -319,7 +322,8 @@ func (rl *ReadLoop) associatePBQAndPnF(ctx context.Context, partitionID partitio
 		// since we created a brand new PBQ it means there is no PnF listening on this PBQ.
 		// we should create and attach the read side of the loop (PnF) to the partition and then
 		// start process-and-forward (pnf) loop
-		rl.op.schedulePnF(ctx, rl.UDF, q, partitionID, rl.toBuffers, rl.whereToDecider, rl.publishWatermark)
+		t := rl.op.schedulePnF(ctx, rl.UDF, q, partitionID, rl.toBuffers, rl.whereToDecider, rl.publishWatermark)
+		rl.udfInvocationTracking[partitionID] = t
 		rl.log.Debugw("Successfully Created/Found pbq and started PnF", zap.String("partitionID", partitionID.String()))
 	}
 
@@ -362,6 +366,9 @@ func (rl *ReadLoop) closePartitions(partitions []partition.ID) {
 	for _, p := range partitions {
 		q := rl.pbqManager.GetPBQ(p)
 		rl.log.Infow("Close of book", zap.String("partitionID", p.String()))
+		// schedule the task for ordered processing.
+		rl.op.insertTask(rl.udfInvocationTracking[p])
 		q.CloseOfBook()
+		delete(rl.udfInvocationTracking, p)
 	}
 }


Test patch:
diff --git a/pkg/reduce/readloop/ordered_test.go b/pkg/reduce/readloop/ordered_test.go
index 0bb3616c..188ae9d4 100644
--- a/pkg/reduce/readloop/ordered_test.go
+++ b/pkg/reduce/readloop/ordered_test.go
@@ -113,11 +113,12 @@ func TestOrderedProcessing(t *testing.T) {
 			defer cancelFn()
 			for _, _partition := range tt.partitions {
 				p, _ := pbqManager.CreateNewPBQ(ctx, _partition)
-				op.schedulePnF(cCtx, identityReducer, p, _partition, toSteps, myForwardTest{}, pw)
+				t := op.schedulePnF(cCtx, identityReducer, p, _partition, toSteps, myForwardTest{}, pw)
+				op.insertTask(t)
 			}
 			assert.Equal(t, op.taskQueue.Len(), tt.expectedBefore)
 			count := 0
-			for e := op.taskQueue.Back(); e != nil; e = e.Prev() {
+			for e := op.taskQueue.Front(); e != nil; e = e.Next() {
 				pfTask := e.Value.(*task)
 				assert.Equal(t, partition.ID{Key: fmt.Sprintf("partition-%d", count)}, pfTask.pf.PartitionID)
 				count = count + 1