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

   1 package main
   2
   3 import (
   4         "container/heap"
   5         "slices"
   6 )
   7
   8 const (
   9         // maxBuilds is the number of model builds that can run at once.
  10         maxBuilds = 8
  11
  12         // leafWindow is the number of unbuilt leaves that can be before a leaf
  13         // whose build starts.
  14         leafWindow = 8
  15 )
  16
  17 // queue holds the nodes that wait for a model build. A leaf joins it when its
  18 // 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 leafWindow leaves before it are
  27 // unbuilt. A failed build waits until the next message.
  28 type queue struct {
  29         leaves  []int
  30         merges  mergeHeap
  31         queued  map[nodeKey]bool
  32         running map[nodeKey]bool
  33         failed  map[nodeKey]bool
  34         retry   []nodeKey
  35 }
  36
  37 // mergeHeap holds the merges that can start, by the end of their context.
  38 type mergeHeap []nodeKey
  39
  40 func newQueue() *queue {
  41         return &queue{
  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 *queue) 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 *queue) 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 *queue) Next() (nodeKey, bool) {
  66         if len(q.running) >= maxBuilds {
  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 // nextLeaf returns the oldest leaf in the window that can start.
  86 func (q *queue) nextLeaf() (nodeKey, bool) {
  87         for _, id := range q.leaves[:min(len(q.leaves), leafWindow)] {
  88                 key := nodeKey{0, id}
  89                 if !q.running[key] && !q.failed[key] {
  90                         return key, true
  91                 }
  92         }
  93
  94         return nodeKey{}, false
  95 }
  96
  97 // Done removes a built node from the queue.
  98 func (q *queue) Done(key nodeKey) {
  99         delete(q.running, key)
 100         if key.level > 0 {
 101                 delete(q.queued, key)
 102                 return
 103         }
 104
 105         position := slices.Index(q.leaves, key.index)
 106         if position >= 0 {
 107                 q.leaves = slices.Delete(q.leaves, position, position+1)
 108         }
 109 }
 110
 111 // Fail keeps a node in the queue, but it cannot start until Retry.
 112 func (q *queue) Fail(key nodeKey) {
 113         delete(q.running, key)
 114         if key.level > 0 {
 115                 q.retry = append(q.retry, key)
 116                 return
 117         }
 118
 119         q.failed[key] = true
 120 }
 121
 122 // Retry lets the failed nodes start again.
 123 func (q *queue) Retry() {
 124         clear(q.failed)
 125         for _, key := range q.retry {
 126                 heap.Push(&q.merges, key)
 127         }
 128
 129         q.retry = q.retry[:0]
 130 }
 131
 132 // Running returns the number of builds that run.
 133 func (q *queue) Running() int {
 134         return len(q.running)
 135 }
 136
 137 func (h mergeHeap) Len() int {
 138         return len(h)
 139 }
 140
 141 func (h mergeHeap) 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 mergeHeap) Swap(a, b int) {
 146         h[a], h[b] = h[b], h[a]
 147 }
 148
 149 func (h *mergeHeap) Push(value any) {
 150         *h = append(*h, value.(nodeKey))
 151 }
 152
 153 func (h *mergeHeap) Pop() any {
 154         old := *h
 155         key := old[len(old)-1]
 156         *h = old[:len(old)-1]
 157
 158         return key
 159 }