chore(engine): Fix calculation of v2 engine range in middleware (#19830)

Author: chaudum

pkg/loki/modules.go

@@ -1181,7 +1181,7 @@ func (t *Loki) initQueryFrontendMiddleware() (_ services.Service, err error) {
 
 		v2Router = queryrange.RouterConfig{
 			Start: start,
-			End:   end,
+			Lag:   t.Cfg.Querier.DataobjStorageLag,
 
 			Validate: engine_v2.IsQuerySupported,
 			Handler:  handler,

pkg/querier/queryrange/engine_router.go

@@ -7,6 +7,7 @@ import (
 	"slices"
 	"time"
 
+	"github.com/coder/quartz"
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
 	"github.com/pkg/errors"
@@ -20,8 +21,8 @@ import (
 
 // RouterConfig configures sending queries to a separate engine.
 type RouterConfig struct {
-	// Start and End time range supported by the engine.
-	Start, End time.Time
+	Start time.Time     // Start time of the v2 engine
+	Lag   time.Duration // Lag after which v2 engine has data
 
 	// Validate function to check if the query is supported by the engine.
 	Validate func(params logql.Params) bool
@@ -38,7 +39,9 @@ type engineReqResp struct {
 
 // engineRouter handles splitting queries between V1 and V2 engines
 type engineRouter struct {
-	v2Start, v2End time.Time // v2 engine time range
+	v2Start time.Time
+	v2Lag   time.Duration
+
 	forMetricQuery bool
 
 	v1Next queryrangebase.Handler
@@ -49,6 +52,9 @@ type engineRouter struct {
 	merger queryrangebase.Merger
 
 	logger log.Logger
+
+	// Used for tests.
+	clock quartz.Clock
 }
 
 // newEngineRouterMiddleware creates a middleware that splits and routes part of the query
@@ -67,21 +73,23 @@ func newEngineRouterMiddleware(
 	return queryrangebase.MiddlewareFunc(func(next queryrangebase.Handler) queryrangebase.Handler {
 		return &engineRouter{
 			v2Start:        v2Config.Start,
-			v2End:          v2Config.End,
+			v2Lag:          v2Config.Lag,
 			v1Next:         queryrangebase.MergeMiddlewares(v1Chain...).Wrap(next),
 			v2Next:         v2Config.Handler,
 			checkV2:        v2Config.Validate,
 			merger:         merger,
 			logger:         logger,
 			forMetricQuery: metricQuery,
+			clock:          quartz.NewReal(),
 		}
 	})
 }
 
 func (e *engineRouter) Do(ctx context.Context, r queryrangebase.Request) (queryrangebase.Response, error) {
+	start, end := e.v2Start, e.getEnd()
 	// if query is entirely before or after v2 engine range, process using next handler.
 	// ignore any boundary overlap, splitting requests that fall on bounary would result in tiny requests.
-	if !r.GetEnd().After(e.v2Start) || !r.GetStart().Before(e.v2End) {
+	if !e.isOverlappingV2range(r, start, end) {
 		return e.v1Next.Do(ctx, r)
 	}
 
@@ -95,7 +103,7 @@ func (e *engineRouter) Do(ctx context.Context, r queryrangebase.Request) (queryr
 		return e.v1Next.Do(ctx, r)
 	}
 
-	inputs := e.splitOverlapping(r, e.v2Start, e.v2End)
+	inputs := e.splitOverlapping(r, start, end)
 
 	// for log queries, order the splits to return early on hitting limits.
 	var limit uint32
@@ -128,6 +136,19 @@ func (e *engineRouter) Do(ctx context.Context, r queryrangebase.Request) (queryr
 	return e.merger.MergeResponse(responses...)
 }
 
+// whether the time range of the request overlaps with the time range of the v2 engine
+func (e engineRouter) isOverlappingV2range(r queryrangebase.Request, start, end time.Time) bool {
+	if !r.GetEnd().After(start) || !r.GetStart().Before(end) {
+		return false
+	}
+	return true
+}
+
+// the end time of the v2 engine based on current timestamp and v2 engine lag
+func (e engineRouter) getEnd() time.Time {
+	return e.clock.Now().UTC().Add(-e.v2Lag)
+}
+
 // splitOverlapping breaks down the request into multiple ranges based on the V2 engine time range.
 // It returns a max of 3 requests:
 // - one for the range before V2 engine

pkg/querier/queryrange/engine_router_test.go

@@ -6,6 +6,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/coder/quartz"
 	"github.com/go-kit/log"
 	"github.com/stretchr/testify/require"
 
@@ -284,10 +285,12 @@ func Test_engineRouter_Do(t *testing.T) {
 		return buildReponse(r, "new"), nil
 	})
 
-	now := time.Now().Truncate(time.Second)
+	clock := quartz.NewMock(t)
+	now := clock.Now().Truncate(time.Second)
+
 	routerConfig := RouterConfig{
 		Start:    now.Add(-24 * time.Hour),
-		End:      now.Add(-time.Hour),
+		Lag:      time.Hour,
 		Validate: func(_ logql.Params) bool { return true },
 		Handler:  v2EngineHandler,
 	}
@@ -300,6 +303,9 @@ func Test_engineRouter_Do(t *testing.T) {
 		log.NewNopLogger(),
 	).Wrap(next)
 
+	routerImpl := router.(*engineRouter)
+	routerImpl.clock = clock
+
 	tests := []struct {
 		name string
 		req  *LokiRequest