From 26f8e2ce1a539acb366239f59aee9df425c7d665 Mon Sep 17 00:00:00 2001
From: Alexandre Bourget <alex@streamingfast.io>
Date: Wed, 5 Jul 2023 02:31:35 -0400
Subject: [PATCH] Fix latest bugs.

---
 orchestrator/stage/stages_test.go | 94 +++++++++++++------------------
 storage/execout/writer.go         |  2 -
 storage/execout/writer_test.go    |  2 +-
 3 files changed, 41 insertions(+), 57 deletions(-)

diff --git a/orchestrator/stage/stages_test.go b/orchestrator/stage/stages_test.go
index 21d111e6c..b0a32265f 100644
--- a/orchestrator/stage/stages_test.go
+++ b/orchestrator/stage/stages_test.go
@@ -24,43 +24,23 @@ func TestNewStages(t *testing.T) {
 		"trace",
 	)
 
-	assert.Equal(t, 8, stages.storeSegmenter.Count()) // from 5 to 75
-	assert.Equal(t, false, stages.storeSegmenter.EndsOnInterval(7))
+	assert.Equal(t, 8, stages.globalSegmenter.Count()) // from 5 to 75
+	assert.Equal(t, true, stages.storeSegmenter.EndsOnInterval(6))
+	assert.Equal(t, false, stages.globalSegmenter.EndsOnInterval(7))
 	assert.Equal(t, 6, stages.storeSegmenter.IndexForStartBlock(60), "index in range")
 	assert.Equal(t, 8, stages.storeSegmenter.IndexForStartBlock(80), "index out of range still returned here")
 	assert.Nil(t, stages.storeSegmenter.Range(8), "out of range")
 
 	assert.Equal(t, block.ParseRange("5-10"), stages.storeSegmenter.Range(0))
 	assert.Equal(t, block.ParseRange("10-20"), stages.storeSegmenter.Range(1))
-	assert.Equal(t, block.ParseRange("70-75"), stages.storeSegmenter.Range(7))
-}
-
-func TestNewStagesMapNotAlignedWithStoreEndBlock(t *testing.T) {
-	reqPlan := plan.BuildTier1RequestPlan(true, 10, 5, 5, 75, 75)
-	assert.Equal(t, "", reqPlan.String())
-
-	stages := NewStages(
-		context.Background(),
-		outputmodules.TestGraphStagedModules(5, 7, 12, 22, 25),
-		reqPlan,
-		nil,
-		"trace",
-	)
-
-	assert.Equal(t, 8, stages.storeSegmenter.Count()) // from 5 to 75
-	assert.Equal(t, false, stages.storeSegmenter.EndsOnInterval(7))
-	assert.Equal(t, 6, stages.storeSegmenter.IndexForStartBlock(60), "index in range")
-	assert.Equal(t, 8, stages.storeSegmenter.IndexForStartBlock(80), "index out of range still returned here")
-	assert.Nil(t, stages.storeSegmenter.Range(8), "out of range")
-
-	assert.Equal(t, block.ParseRange("5-10"), stages.storeSegmenter.Range(0))
-	assert.Equal(t, block.ParseRange("10-20"), stages.storeSegmenter.Range(1))
-	assert.Equal(t, block.ParseRange("70-75"), stages.storeSegmenter.Range(7))
+	assert.Nil(t, stages.storeSegmenter.Range(7))
+	assert.Equal(t, block.ParseRange("70-75"), stages.globalSegmenter.Range(7))
 }
 
 func TestNewStagesNextJobs(t *testing.T) {
 	//seg := block.NewSegmenter(10, 5, 50)
 	reqPlan := plan.BuildTier1RequestPlan(true, 10, 5, 5, 50, 50)
+	assert.Equal(t, "interval=10, stores=[5, 40), map_write=[5, 50), map_read=[5, 50), linear=[nil)", reqPlan.String())
 	stages := NewStages(
 		context.Background(),
 		outputmodules.TestGraphStagedModules(5, 5, 5, 5, 5),
@@ -159,61 +139,67 @@ S:CCSS
 S:CSS.
 M:NC..`)
 
-	stages.NextJob()
-
-	segmentStateEquals(t, stages, `
-S:CCSSS...
-S:CSS.....
-M:NC......`)
-
-	stages.NextJob()
+	_, r := stages.NextJob()
+	assert.Nil(t, r)
 
 	segmentStateEquals(t, stages, `
-S:CCSSSS..
-S:CSS.....
-M:NC......`)
+S:CCSS
+S:CSS.
+M:NC..`)
 
-	_, r := stages.NextJob()
-	assert.Nil(t, r)
+	//	segmentStateEquals(t, stages, `
+	//S:CCSSS...
+	//S:CSS.....
+	//M:NC......`)
+	//
+	//	stages.NextJob()
+	//
+	//	segmentStateEquals(t, stages, `
+	//S:CCSSSS..
+	//S:CSS.....
+	//M:NC......`)
+	//
+	//	_, r := stages.NextJob()
+	//	assert.Nil(t, r)
 	stages.MarkSegmentPartialPresent(id(2, 0))
 
 	segmentStateEquals(t, stages, `
-S:CCPSSS..
-S:CSS.....
-M:NC......`)
+S:CCPS
+S:CSS.
+M:NC..`)
 
 	_, r = stages.NextJob()
 	assert.Nil(t, r)
 	stages.MarkSegmentMerging(id(2, 0))
 
 	segmentStateEquals(t, stages, `
-S:CCMSSS..
-S:CSS.....
-M:NC......`)
+S:CCMS
+S:CSS.
+M:NC..`)
 
 	_, r = stages.NextJob()
 	assert.Nil(t, r)
 	stages.markSegmentCompleted(id(2, 0))
 
 	segmentStateEquals(t, stages, `
-S:CCCSSS..
-S:CSS.....
-M:NC......`)
+S:CCCS
+S:CSS.
+M:NC..`)
 
 	stages.NextJob()
 
 	segmentStateEquals(t, stages, `
-S:CCCSSS..
-S:CSSS....
-M:NC......`)
+S:CCCS
+S:CSSS
+M:NC..`)
 
 	stages.forceTransition(1, 1, UnitCompleted)
 	stages.NextJob()
 
 	segmentStateEquals(t, stages, `
-S:CCCSSS..
-S:CCSS....
-M:NCS.....`)
+S:CCCS
+S:CCSS
+M:NCS.`)
 
 }
 
diff --git a/storage/execout/writer.go b/storage/execout/writer.go
index a941f40b8..c0a979a3e 100644
--- a/storage/execout/writer.go
+++ b/storage/execout/writer.go
@@ -21,14 +21,12 @@ type Writer struct {
 
 	fileWalker   *FileWalker
 	currentFile  *File
-	files        map[string]*File // moduleName => file
 	outputModule string
 }
 
 func NewWriter(initialBlockBoundary, exclusiveEndBlock uint64, outputModule string, configs *Configs) *Writer {
 	w := &Writer{
 		wg:           &sync.WaitGroup{},
-		files:        make(map[string]*File),
 		outputModule: outputModule,
 	}
 
diff --git a/storage/execout/writer_test.go b/storage/execout/writer_test.go
index 80072737f..7f417648a 100644
--- a/storage/execout/writer_test.go
+++ b/storage/execout/writer_test.go
@@ -25,5 +25,5 @@ var testConfigs = &Configs{
 func TestNewExecOutputWriterIsSubRequest(t *testing.T) {
 	res := NewWriter(11, 15, "A", testConfigs)
 	require.NotNil(t, res)
-	assert.Equal(t, 15, int(res.files["A"].ExclusiveEndBlock))
+	assert.Equal(t, 15, int(res.fileWalker.segmenter.ExclusiveEndBlock()))
 }