From 4df45566a637a21ffc4f92483d577ef068eae599 Mon Sep 17 00:00:00 2001
From: Rob Skillington <rob.skillington@gmail.com>
Date: Fri, 7 May 2021 02:24:26 -0400
Subject: [PATCH] [query] Fix Graphite context time window expansions not being
 cumulative

---
 src/query/graphite/common/context.go          |  8 +--
 src/query/graphite/common/test_util.go        | 16 +++--
 .../graphite/native/builtin_functions_test.go | 67 +++++++++++++++++++
 3 files changed, 80 insertions(+), 11 deletions(-)

diff --git a/src/query/graphite/common/context.go b/src/query/graphite/common/context.go
index 7750925acb..32dc397248 100644
--- a/src/query/graphite/common/context.go
+++ b/src/query/graphite/common/context.go
@@ -161,12 +161,8 @@ func (c *Context) NewChildContext(opts ChildContextOptions) *Context {
 		child.TimeRangeAdjustment.OriginalEnd = origEnd
 		child.TimeRangeAdjustment.ShiftStart += opts.adjustment.shiftStart
 		child.TimeRangeAdjustment.ShiftEnd += opts.adjustment.shiftEnd
-		if opts.adjustment.expandStart > child.TimeRangeAdjustment.ExpandStart {
-			child.TimeRangeAdjustment.ExpandStart = opts.adjustment.expandStart
-		}
-		if opts.adjustment.expandEnd > child.TimeRangeAdjustment.ExpandEnd {
-			child.TimeRangeAdjustment.ExpandEnd = opts.adjustment.expandEnd
-		}
+		child.TimeRangeAdjustment.ExpandStart += opts.adjustment.expandStart
+		child.TimeRangeAdjustment.ExpandEnd += opts.adjustment.expandEnd
 
 		child.StartTime = origStart.
 			Add(child.TimeRangeAdjustment.ShiftStart).
diff --git a/src/query/graphite/common/test_util.go b/src/query/graphite/common/test_util.go
index 4e832f0166..a41a322c5c 100644
--- a/src/query/graphite/common/test_util.go
+++ b/src/query/graphite/common/test_util.go
@@ -92,16 +92,22 @@ func NewConsolidationTestSeries(start, end time.Time, duration time.Duration) (*
 }
 
 // CompareOutputsAndExpected compares the actual output with the expected output.
-func CompareOutputsAndExpected(t *testing.T, step int, start time.Time, expected []TestSeries,
-	actual []*ts.Series) {
-	require.Equal(t, len(expected), len(actual))
+func CompareOutputsAndExpected(
+	t *testing.T,
+	step int,
+	start time.Time,
+	expected []TestSeries,
+	actual []*ts.Series,
+) {
+	require.Equal(t, len(expected), len(actual), "mismatch series count")
 	for i := range expected {
 		i := i // To capture for wrapMsg.
 		e := expected[i].Data
 		a := actual[i]
 		wrapMsg := func(str string) string {
-			return fmt.Sprintf("%s\nseries=%d\nexpected=%v\nactual=%v",
-				str, i, e, a.SafeValues())
+			return fmt.Sprintf("%s\nseries=%d\nexpected=%v\nactual=%v\n"+
+				"expectedStart=%v\nactualStart=%v\n",
+				str, i, e, a.SafeValues(), start, a.StartTime())
 		}
 		require.Equal(t, expected[i].Name, a.Name())
 		assert.Equal(t, step, a.MillisPerStep(), wrapMsg(a.Name()+
diff --git a/src/query/graphite/native/builtin_functions_test.go b/src/query/graphite/native/builtin_functions_test.go
index 60c46866b1..2332c05a2f 100644
--- a/src/query/graphite/native/builtin_functions_test.go
+++ b/src/query/graphite/native/builtin_functions_test.go
@@ -1046,6 +1046,73 @@ func TestMovingSumSuccess(t *testing.T) {
 	testMovingFunction(t, "movingSum(foo.bar.baz, '30s')", "movingSum(foo.bar.baz,3)", nil, nil, nil)
 }
 
+// TestMovingSumOfMovingSum tests that expansion of the time window
+// fetched is stacked when child contexts are created, otherwise the child
+// functions do not have enough data to work with to satisfy the parent call.
+func TestMovingSumOfMovingSum(t *testing.T) {
+	var (
+		ctrl   = xgomock.NewController(t)
+		store  = storage.NewMockStorage(ctrl)
+		engine = NewEngine(store, CompileOptions{})
+		end    = time.Now().Truncate(time.Minute)
+		start  = end.Add(-5 * time.Minute)
+		ctx    = common.NewContext(common.ContextOptions{
+			Start:  start,
+			End:    end,
+			Engine: engine,
+		})
+		millisPerStep = 60000
+		data          = []float64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}
+	)
+
+	defer ctrl.Finish()
+	defer func() { _ = ctx.Close() }()
+
+	store.EXPECT().
+		FetchByQuery(gomock.Any(), "foo.bar", gomock.Any()).
+		DoAndReturn(func(
+			_ *common.Context,
+			query string,
+			opts storage.FetchOptions,
+		) (*storage.FetchResult, error) {
+			if !opts.EndTime.Equal(end) {
+				return nil, fmt.Errorf("unexpected end")
+			}
+			length := int(opts.EndTime.Sub(opts.StartTime) / time.Minute)
+			values := data[len(data)-length:]
+			return &storage.FetchResult{
+				SeriesList: []*ts.Series{
+					ts.NewSeries(ctx, query, opts.StartTime,
+						common.NewTestSeriesValues(ctx, millisPerStep, values)),
+				},
+			}, nil
+		}).
+		AnyTimes()
+
+	target := `movingSum(movingSum(foo.bar,"2min"),"5min")`
+
+	phonyContext := common.NewContext(common.ContextOptions{
+		Start:  start,
+		End:    end,
+		Engine: engine,
+	})
+
+	expr, err := phonyContext.Engine.(*Engine).Compile(target)
+	require.NoError(t, err)
+	res, err := expr.Execute(phonyContext)
+	require.NoError(t, err)
+
+	expected := []common.TestSeries{
+		{
+			Name: target,
+			Data: []float64{65, 75, 85, 95, 105},
+		},
+	}
+
+	common.CompareOutputsAndExpected(t, millisPerStep, start,
+		expected, res.Values)
+}
+
 func TestMovingSumError(t *testing.T) {
 	testMovingFunctionError(t, "movingSum(foo.bar.baz, '-30s')")
 	testMovingFunctionError(t, "movingSum(foo.bar.baz, 0)")