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         if key.level > 0 {
  86                 delete(q.queued, key)
  87                 return
  88         }
  89
  90         position := slices.Index(q.leaves, key.index)
  91         if position >= 0 {
  92                 q.leaves = slices.Delete(q.leaves, position, position+1)
  93         }
  94 }
  95
  96 // Fail keeps a node in the queue, but it cannot start until Retry.
  97 func (q *BuildQueue) Fail(key NodeKey) {
  98         delete(q.running, key)
  99         if key.level > 0 {
 100                 q.retry = append(q.retry, key)
 101                 return
 102         }
 103
 104         q.failed[key] = true
 105 }
 106
 107 // Retry lets the failed nodes start again.
 108 func (q *BuildQueue) Retry() {
 109         clear(q.failed)
 110         for _, key := range q.retry {
 111                 heap.Push(&q.merges, key)
 112         }
 113
 114         q.retry = q.retry[:0]
 115 }
 116
 117 // Running returns the number of builds that run.
 118 func (q *BuildQueue) Running() int {
 119         return len(q.running)
 120 }
 121
 122 // nextLeaf returns the oldest leaf in the window that can start.
 123 func (q *BuildQueue) nextLeaf() (NodeKey, bool) {
 124         for _, id := range q.leaves[:min(len(q.leaves), buildLeafWindow)] {
 125                 key := NodeKey{0, id}
 126                 if !q.running[key] && !q.failed[key] {
 127                         return key, true
 128                 }
 129         }
 130
 131         return NodeKey{}, false
 132 }
 133
 134 // BuildMergeHeap holds the merges that can start, by the end of their context.
 135 type BuildMergeHeap []NodeKey
 136
 137 func (h BuildMergeHeap) Len() int {
 138         return len(h)
 139 }
 140
 141 func (h BuildMergeHeap) Less(a, b int) bool {
 142         return nodeEnd(h[a].level, h[a].index) < nodeEnd(h[b].level, h[b].index)
 143 }
 144
 145 func (h BuildMergeHeap) Swap(a, b int) {
 146         h[a], h[b] = h[b], h[a]
 147 }
 148
 149 func (h *BuildMergeHeap) Push(value any) {
 150         *h = append(*h, value.(NodeKey))
 151 }
 152
 153 func (h *BuildMergeHeap) Pop() any {
 154         old := *h
 155         key := old[len(old)-1]
 156         *h = old[:len(old)-1]
 157
 158         return key
 159 }