From 6928aac776ff6612464908c9c7d452966d70773a Mon Sep 17 00:00:00 2001
From: Bryan Boreham <bjboreham@gmail.com>
Date: Fri, 23 Jun 2023 21:48:55 +0100
Subject: [PATCH] loser tree: allow sequences to contain maxVal (#9057)

**What this PR does / why we need it**:

The tree stores the loser in each node, or maxVal if the sequence has
ended. If sequences can contain maxVal, we need to take care not to let
an ended sequence overtake an active one.

Some test cases from Charles who pointed out the issue in #9053.

Note this situation cannot happen for real in Loki, since we don't
expect log lines with a timestamp that far into the future,
but this is supposed to be a re-usable library.

**Checklist**
- [x] Reviewed the
[`CONTRIBUTING.md`](https://github.com/grafana/loki/blob/main/CONTRIBUTING.md)
guide (**required**)
- NA Documentation added
- [x] Tests updated
- NA `CHANGELOG.md` updated
- NA Changes that require user attention or interaction to upgrade are
documented in `docs/sources/upgrading/_index.md`

Co-authored-by: Charles Korn <charles.korn@grafana.com>
---
 pkg/util/loser/tree.go      | 33 +++++++++++++++++++++++++++++----
 pkg/util/loser/tree_test.go | 25 +++++++++++++++++++++++++
 2 files changed, 54 insertions(+), 4 deletions(-)

diff --git a/pkg/util/loser/tree.go b/pkg/util/loser/tree.go
index 3810622ab5..aa99e02991 100644
--- a/pkg/util/loser/tree.go
+++ b/pkg/util/loser/tree.go
@@ -79,8 +79,11 @@ func (t *Tree[E, S]) Next() bool {
 	if t.nodes[t.nodes[0].index].index == -1 { // already exhausted
 		return false
 	}
-	t.moveNext(t.nodes[0].index)
-	t.replayGames(t.nodes[0].index)
+	if t.moveNext(t.nodes[0].index) {
+		t.replayGames(t.nodes[0].index)
+	} else {
+		t.sequenceEnded(t.nodes[0].index)
+	}
 	return t.nodes[t.nodes[0].index].index != -1
 }
 
@@ -102,11 +105,12 @@ func (t *Tree[E, S]) initialize() {
 	t.nodes[0].value = t.nodes[winners[1]].value
 }
 
-// Starting at pos, re-consider all values up to the root.
+// Starting at pos, which is a winner, re-consider all values up to the root.
 func (t *Tree[E, S]) replayGames(pos int) {
-	// At the start, pos is a leaf node, and is the winner at that level.
 	n := parent(pos)
 	for n != 0 {
+		// If n.value < pos.value then pos loses.
+		// If they are equal, pos wins because n could be a sequence that ended, with value maxval.
 		if t.less(t.nodes[n].value, t.nodes[pos].value) {
 			loser := pos
 			// Record pos as the loser here, and the old loser is the new winner.
@@ -121,6 +125,27 @@ func (t *Tree[E, S]) replayGames(pos int) {
 	t.nodes[0].value = t.nodes[pos].value
 }
 
+func (t *Tree[E, S]) sequenceEnded(pos int) {
+	// Find the first active sequence which used to lose to it.
+	n := parent(pos)
+	for n != 0 && t.nodes[t.nodes[n].index].index == -1 {
+		n = parent(n)
+	}
+	if n == 0 {
+		// There are no active sequences left
+		t.nodes[0].index = pos
+		t.nodes[0].value = t.maxVal
+		return
+	}
+
+	// Record pos as the loser here, and the old loser is the new winner.
+	loser := pos
+	winner := t.nodes[n].index
+	t.nodes[n].index = loser
+	t.nodes[n].value = t.nodes[loser].value
+	t.replayGames(winner)
+}
+
 func (t *Tree[E, S]) playGame(a, b int) (loser, winner int) {
 	if t.less(t.nodes[a].value, t.nodes[b].value) {
 		return b, a
diff --git a/pkg/util/loser/tree_test.go b/pkg/util/loser/tree_test.go
index e71610aa41..956b420f12 100644
--- a/pkg/util/loser/tree_test.go
+++ b/pkg/util/loser/tree_test.go
@@ -94,6 +94,31 @@ var testCases = []struct {
 		args: []*List{NewList(1, 3), NewList(2, 4, 5)},
 		want: NewList(1, 2, 3, 4, 5),
 	},
+	{
+		name: "two lists, largest value in first list equal to maximum",
+		args: []*List{NewList(1, math.MaxUint64), NewList(2, 3)},
+		want: NewList(1, 2, 3, math.MaxUint64),
+	},
+	{
+		name: "two lists, first straddles second and has maxval",
+		args: []*List{NewList(1, 3, math.MaxUint64), NewList(2)},
+		want: NewList(1, 2, 3, math.MaxUint64),
+	},
+	{
+		name: "two lists, largest value in second list equal to maximum",
+		args: []*List{NewList(1, 3), NewList(2, math.MaxUint64)},
+		want: NewList(1, 2, 3, math.MaxUint64),
+	},
+	{
+		name: "two lists, second straddles first and has maxval",
+		args: []*List{NewList(2), NewList(1, 3, math.MaxUint64)},
+		want: NewList(1, 2, 3, math.MaxUint64),
+	},
+	{
+		name: "two lists, largest value in both lists equal to maximum",
+		args: []*List{NewList(1, math.MaxUint64), NewList(2, math.MaxUint64)},
+		want: NewList(1, 2, math.MaxUint64, math.MaxUint64),
+	},
 	{
 		name: "three lists",
 		args: []*List{NewList(1, 3), NewList(2, 4), NewList(5)},