git.giacomolli.uk - onechat/blob - build_queue.go

   1 package main
   2
   3 import (
   4         "container/heap"
   5         "slices"
   6 )
   7
   8 const (
   9         // buildMax is the number of model builds that can run at once.
  10         buildMax = 8
  11
  12         // buildLeafWindow is the number of unbuilt leaves that can be before a
  13         // leaf whose build starts.
  14         buildLeafWindow = 8
  15 )
  16
  17 // BuildQueue holds the nodes that wait for a model build. A leaf joins it when
  18 // its message is too long to be its own node. A merge joins it when both of its
  19 // children are built and their texts do not fit in one node. So the core never
  20 // scans the tree for work.
  21 //
  22 // The builds start in the order of the messages: the node whose context ends
  23 // first starts first. The context of leaf i ends before message i, and the
  24 // context of a merge ends after its last message. At a tie the merge starts
  25 // first. So a backlog of leaves does not stop the merges that the view needs to
  26 // shrink. A leaf can start only when fewer than buildLeafWindow leaves before
  27 // it are unbuilt. A failed build waits until the next message.
  28 type BuildQueue struct {
  29         leaves  []int
  30         merges  BuildMergeHeap
  31         queued  map[NodeKey]bool
  32         running map[NodeKey]bool
  33         failed  map[NodeKey]bool
  34         retry   []NodeKey
  35 }
  36
  37 // BuildMergeHeap holds the merges that can start, by the end of their context.
  38 type BuildMergeHeap []NodeKey
  39
  40 func newBuildQueue() *BuildQueue {
  41         return &BuildQueue{
  42                 queued:  make(map[NodeKey]bool),
  43                 running: make(map[NodeKey]bool),
  44                 failed:  make(map[NodeKey]bool),
  45         }
  46 }
  47
  48 // AddLeaf adds the leaf of message id. The IDs must come in increasing order.
  49 func (q *BuildQueue) AddLeaf(id int) {
  50         q.leaves = append(q.leaves, id)
  51 }
  52
  53 // AddMerge adds a merge whose children are built. A merge that is in the queue
  54 // already stays as it is.
  55 func (q *BuildQueue) AddMerge(key NodeKey) {
  56         if q.queued[key] {
  57                 return
  58         }
  59
  60         q.queued[key] = true
  61         heap.Push(&q.merges, key)
  62 }
  63
  64 // Next returns the next node that can start, and marks it as running.
  65 func (q *BuildQueue) Next() (NodeKey, bool) {
  66         if len(q.running) >= buildMax {
  67                 return NodeKey{}, false
  68         }
  69
  70         leaf, found := q.nextLeaf()
  71         if len(q.merges) > 0 && (!found || nodeEnd(q.merges[0].level, q.merges[0].index) <= leaf.index) {
  72                 key := heap.Pop(&q.merges).(NodeKey)
  73                 q.running[key] = true
  74
  75                 return key, true
  76         }
  77
  78         if found {
  79                 q.running[leaf] = true
  80         }
  81
  82         return leaf, found
  83 }
  84
  85 // Done removes a built node from the queue.
  86 func (q *BuildQueue) Done(key NodeKey) {
  87         delete(q.running, key)
  88
  89         if key.level > 0 {
  90                 delete(q.queued, key)
  91
  92                 return
  93         }
  94
  95         position := slices.Index(q.leaves, key.index)
  96         if position >= 0 {
  97                 q.leaves = slices.Delete(q.leaves, position, position+1)
  98         }
  99 }
 100
 101 // Fail keeps a node in the queue, but it cannot start until Retry.
 102 func (q *BuildQueue) Fail(key NodeKey) {
 103         delete(q.running, key)
 104
 105         if key.level > 0 {
 106                 q.retry = append(q.retry, key)
 107                 return
 108         }
 109
 110         q.failed[key] = true
 111 }
 112
 113 // Retry lets the failed nodes start again.
 114 func (q *BuildQueue) Retry() {
 115         clear(q.failed)
 116
 117         for _, key := range q.retry {
 118                 heap.Push(&q.merges, key)
 119         }
 120
 121         q.retry = q.retry[:0]
 122 }
 123
 124 // Running returns the number of builds that run.
 125 func (q *BuildQueue) Running() int {
 126         return len(q.running)
 127 }
 128
 129 // nextLeaf returns the oldest leaf in the window that can start.
 130 func (q *BuildQueue) nextLeaf() (NodeKey, bool) {
 131         for _, id := range q.leaves[:min(len(q.leaves), buildLeafWindow)] {
 132                 key := NodeKey{0, id}
 133                 if !q.running[key] && !q.failed[key] {
 134                         return key, true
 135                 }
 136         }
 137
 138         return NodeKey{}, false
 139 }
 140
 141 func (h BuildMergeHeap) Len() int {
 142         return len(h)
 143 }
 144
 145 func (h BuildMergeHeap) Less(a, b int) bool {
 146         return nodeEnd(h[a].level, h[a].index) < nodeEnd(h[b].level, h[b].index)
 147 }
 148
 149 func (h BuildMergeHeap) Swap(a, b int) {
 150         h[a], h[b] = h[b], h[a]
 151 }
 152
 153 func (h *BuildMergeHeap) Push(value any) {
 154         *h = append(*h, value.(NodeKey))
 155 }
 156
 157 func (h *BuildMergeHeap) Pop() any {
 158         old := *h
 159         key := old[len(old)-1]
 160         *h = old[:len(old)-1]
 161
 162         return key
 163 }