From e2e7cee7c2e64ed9cf72d5dfc9ecad24f9e4cb17 Mon Sep 17 00:00:00 2001
From: Thomas Jackson <thomas@conduit.xyz>
Date: Tue, 27 Aug 2024 22:09:09 -0700
Subject: [PATCH] Add NodeReplacer for BinaryExpr where it is a VectorSelector
 + Literal

If the query is something like `foo>1` prior to this patch we unwrap and
send down to each downstream `foo` and then do the comparison
(binaryExpr) in promxy. This is non-ideal as we are sending some data
back that could have been dropped at the remote end.

Although this is a great performance optimization we will need to be
careful when adding cases to this. For BinaryExprs to be sent to
each servergroup in a `valid` way the query must be "local" to a given
servergroup. For example; `sum(foo)>1` would *not* be safe to send
downstream as `sum(foo)` requires adding across a variety of
servergroups to determine the correct value.

Fixes #673
---
 pkg/proxystorage/proxy.go | 78 +++++++++++++++++++++++++++++++++++++++
 pkg/proxystorage/util.go  |  8 ++++
 2 files changed, 86 insertions(+)

diff --git a/pkg/proxystorage/proxy.go b/pkg/proxystorage/proxy.go
index 7027abe30..605a35f0f 100644
--- a/pkg/proxystorage/proxy.go
+++ b/pkg/proxystorage/proxy.go
@@ -733,6 +733,84 @@ func (p *ProxyStorage) NodeReplacer(ctx context.Context, s *parser.EvalStmt, nod
 			return n, nil
 		}
 
+	// BinaryExprs *can* be sent untouched to downstreams assuming there is no actual interaction between LHS/RHS
+	// these are relatively rare -- as things like `sum(foo) > 2` would *not* be viable as `sum(foo)` could
+	// potentially require multiple servergroups to generate the correct response.
+	// For now we'll start a VERY basic check for a `VectorSelector` and either of (`NumberLiteral` or `StringLiteral`)
+	case *parser.BinaryExpr:
+
+		vectorBinaryExpr := func(vs *parser.VectorSelector) (parser.Node, error) {
+			logrus.Debugf("BinaryExpr (VectorSelector + Literal): %v", n)
+			removeOffsetFn()
+
+			var result model.Value
+			var warnings v1.Warnings
+			var err error
+			if s.Interval > 0 {
+				vs.LookbackDelta = s.Interval - time.Duration(1)
+				result, warnings, err = state.client.QueryRange(ctx, n.String(), v1.Range{
+					Start: s.Start.Add(-offset),
+					End:   s.End.Add(-offset),
+					Step:  s.Interval,
+				})
+			} else {
+				result, warnings, err = state.client.Query(ctx, n.String(), s.Start.Add(-offset))
+			}
+
+			if err != nil {
+				return nil, err
+			}
+
+			iterators := promclient.IteratorsForValue(result)
+			series := make([]storage.Series, len(iterators))
+			for i, iterator := range iterators {
+				series[i] = &proxyquerier.Series{iterator}
+			}
+
+			ret := &parser.VectorSelector{OriginalOffset: offset}
+			if s.Interval > 0 {
+				ret.LookbackDelta = s.Interval - time.Duration(1)
+			}
+			ret.UnexpandedSeriesSet = proxyquerier.NewSeriesSet(series, promhttputil.WarningsConvert(warnings), err)
+			return ret, nil
+		}
+
+		var (
+			vectorSelector *parser.VectorSelector
+			otherSide      parser.Expr
+		)
+
+		if vs, ok := n.LHS.(*parser.VectorSelector); ok {
+			vectorSelector = vs
+			otherSide = n.RHS
+		} else if vs, ok := n.RHS.(*parser.VectorSelector); ok {
+			vectorSelector = vs
+			otherSide = n.LHS
+		}
+
+		// If we *have* a vectorselector; we can check the `otherSide`
+		if vectorSelector != nil {
+			// Only valid if the other side is either `NumberLiteral` or `StringLiteral`
+			literal := false
+			switch UnwrapExpr(otherSide).(type) {
+			case *parser.StringLiteral:
+				literal = true
+			case *parser.NumberLiteral:
+				literal = true
+			}
+
+			// If the `otherSide` is a literal; we attempt to unwrap`
+			if literal {
+				ret, err := vectorBinaryExpr(vectorSelector)
+				if err != nil {
+					return nil, err
+				}
+				if ret != nil {
+					return ret, nil
+				}
+			}
+		}
+
 	default:
 		logrus.Debugf("default %v %s", n, reflect.TypeOf(n))
 
diff --git a/pkg/proxystorage/util.go b/pkg/proxystorage/util.go
index 3d1c76f24..2953cbc1a 100644
--- a/pkg/proxystorage/util.go
+++ b/pkg/proxystorage/util.go
@@ -119,3 +119,11 @@ func PreserveLabel(expr parser.Expr, srcLabel string, dstLabel string) (relabelE
 	relabelExpress, _ = parser.ParseExpr(fmt.Sprintf("label_replace(%s,`%s`,`$1`,`%s`,`(.*)`)", expr.String(), dstLabel, srcLabel))
 	return relabelExpress
 }
+
+func UnwrapExpr(expr parser.Expr) parser.Expr {
+	switch e := expr.(type) {
+	case *parser.StepInvariantExpr:
+		return e.Expr
+	}
+	return expr
+}