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