loki/pkg/engine/internal/util/queue/fair/node.go

package fair

import "container/heap"

// node represents an individual element in the hierarchical queue tree.
// Leaf nodes hold values of T, while internal nodes hold [pqueue].
type node[T any] struct {
	parent *node[T]  // Parent this node belongs to.
	queue  *Queue[T] // Queue this node belongs to.

	// Rank of this node within its parent. The rank determines [Queue.Pop]
	// order: lower rank means higher priority.
	rank int64

	// id of this node, assigned when it is created. id is used to tiebreak when
	// two nodes have the same rank (lower id wins).
	id int64

	// value is either a pqueue[T] for internal nodes, or T for leaf elements.
	value any
}

// FixChild fixes the priority of a child after its rank has been modified.
func (n *node[T]) FixChild(name string) {
	pq := n.pqueue()
	child, ok := pq.scopeLookup[name]
	if !ok {
		return
	}
	heap.Fix(pq, child.Index)
}

func (n *node[T]) pqueue() *pqueue[T] {
	pq, ok := n.value.(*pqueue[T])
	if !ok {
		panic("node is not a scope node")
	}
	return pq
}

// Peek returns the next selected child. Panics if n is not a scope node.
func (n *node[T]) Peek() (*node[T], Scope) {
	pq := n.pqueue()
	if len(pq.children) == 0 {
		return nil, nil
	}
	return pq.children[0].node, pq.scope
}

// Pop removes the next selected child from the heap. Panics if n is not a scope node.
func (n *node[T]) Pop() (*node[T], Scope) {
	pq := n.pqueue()
	if len(pq.children) == 0 {
		return nil, nil
	}

	child := pq.children[0].node
	_ = heap.Remove(pq, 0)
	return child, pq.scope
}

// Scope returns the scope for this node. Panics if n is not a scope node.
func (n *node[T]) Scope() Scope { return n.pqueue().scope }

// NumItems returns the number of non-heap children of n.
//
// NumItems panics if n is not a scope node.
func (n *node[T]) NumItems() int {
	pq, ok := n.value.(*pqueue[T])
	if !ok {
		panic("node is not a scope node")
	}

	var total int
	for _, c := range pq.children {
		_, isItem := c.value.(T)
		if isItem {
			total++
		}
	}
	return total
}

// NumChildren returns the number of children scopes of n (alive and dead).
//
// NumChildren panics if n is not a scope node.
func (n *node[T]) NumChildren() int {
	pq := n.pqueue()
	return len(pq.scopeLookup)
}

// Len returns the number of values of n.
//
// Children panics if n is not a scope node.
func (n *node[T]) Len() int { return n.pqueue().Len() }

// MarkAlive flags a child scope as alive, indicating that it has children and
// should be considered for selection in the heap. MarkAlive is a no-op if the
// scope is already alive.
//
// MarkAlive panics if n is not a scope node.
func (n *node[T]) MarkAlive(name string) {
	pq := n.pqueue()
	pn, ok := pq.scopeLookup[name]
	if !ok || pn.Index >= 0 /* already alive */ {
		return
	}

	// Reviving a scope shouldn't jump in line in the priority queue, so we
	// give it a new ID and ensure that its rank is no lower than the minimum
	// rank among its siblings.
	{
		var initRank int64
		if sibling, _ := n.Peek(); sibling != nil {
			initRank = sibling.rank
		}
		pn.rank = max(initRank, pn.rank)
		pn.id = n.queue.getNextID()
	}

	pn.Index = len(pq.children)
	pq.children = append(pq.children, pn)
	heap.Fix(pq, pn.Index)
}

// MarkDead flags the provided scope name as dead (without unregistering it),
// indicating that it has no children and should not be considered for selection
// in the heap. MarkDead is a no-op if the scope is already dead.
//
// MarkDead panics if n is not a scope node.
func (n *node[T]) MarkDead(name string) {
	pq := n.pqueue()
	pn, ok := pq.scopeLookup[name]
	if !ok || pn.Index < 0 /* already dead */ {
		return
	}

	heap.Remove(pq, pn.Index)
	pn.Index = -1 // Flag as dead
}

// GetScope returns the child scope with the given name. Panics if n is not a
// scope node.
func (n *node[T]) GetScope(name string) (*node[T], bool) {
	pq := n.pqueue()
	if n, ok := pq.scopeLookup[name]; ok {
		return n.node, true
	}
	return nil, false
}

// UnregisterScope removes the child scope with the given name.
//
// UnregisterScope panics if n is not a scope node.
func (n *node[T]) UnregisterScope(name string) {
	pq := n.pqueue()
	n.MarkDead(name)             // Remove it from the heap first,
	delete(pq.scopeLookup, name) // then delete the scope.
}

// RegisterScope creates a child scope. Returns an error if the scope already
// exists.
//
// RegisterScope panics if n is not a scope node. The scope is not validated to
// ensure that it is an immediate child of n.
func (n *node[T]) RegisterScope(scope Scope) (*node[T], error) {
	if len(scope) == 0 {
		return nil, ErrEmptyScope
	}

	if _, ok := n.GetScope(scope.Name()); ok {
		return nil, ErrScopeExists
	}

	var initRank int64
	if sibling, _ := n.Peek(); sibling != nil {
		initRank = sibling.rank
	}

	newNode := &node[T]{
		parent: n,
		queue:  n.queue,

		rank: initRank,
		id:   n.queue.getNextID(),

		value: &pqueue[T]{scope: scope},
	}

	// Register the scope. We shouldn't *push* the scope onto the heap yet until
	// it has at least one element in it.
	pq := n.pqueue()
	if pq.scopeLookup == nil {
		pq.scopeLookup = make(map[string]*posNode[T])
	}
	pq.scopeLookup[scope.Name()] = &posNode[T]{
		node:  newNode,
		Index: -1, // Start as inactive
	}
	return newNode, nil
}

// CreateValue pushes a value node to n. If v already exists in n, a duplicate
// value entry for v is added. Panics if n is not a scope node.
func (n *node[T]) CreateValue(v T) *node[T] {
	pq := n.pqueue()

	// The initial rank of a new node is the rank of the first child.
	var initRank int64
	if len(pq.children) > 0 {
		initRank = pq.children[0].rank
	}

	newNode := &node[T]{
		parent: n,
		queue:  n.queue,

		rank: initRank,
		id:   n.queue.getNextID(),

		value: v,
	}

	heap.Push(pq, newNode)
	return newNode
}