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

   1 package main
   2
   3 import (
   4         "slices"
   5 )
   6
   7 const (
   8         // maxBuilds is the number of model builds that can run at once.
   9         maxBuilds = 8
  10
  11         // leafWindow is the number of unbuilt leaves that can be before a leaf
  12         // whose build starts.
  13         leafWindow = 8
  14 )
  15
  16 // queue holds the nodes that wait for a model build. A leaf joins it when its
  17 // message is too long to be its own node. A merge joins it when both of its
  18 // children are built and their texts do not fit in one node. So the core never
  19 // scans the tree for work.
  20 //
  21 // A leaf can start when fewer than leafWindow leaves before it are unbuilt.
  22 // Leaves start before merges, oldest first, because a turn waits for the leaves
  23 // of all earlier messages. Merges start in the order in which they became
  24 // ready. A failed build waits until the next message.
  25 type queue struct {
  26         leaves  []int
  27         merges  []nodeKey
  28         queued  map[nodeKey]bool
  29         running map[nodeKey]bool
  30         failed  map[nodeKey]bool
  31 }
  32
  33 func newQueue() *queue {
  34         return &queue{
  35                 queued:  make(map[nodeKey]bool),
  36                 running: make(map[nodeKey]bool),
  37                 failed:  make(map[nodeKey]bool),
  38         }
  39 }
  40
  41 // AddLeaf adds the leaf of message id. The IDs must come in increasing order.
  42 func (q *queue) AddLeaf(id int) {
  43         q.leaves = append(q.leaves, id)
  44 }
  45
  46 // AddMerge adds a merge whose children are built. A merge that is in the queue
  47 // already stays as it is.
  48 func (q *queue) AddMerge(key nodeKey) {
  49         if q.queued[key] {
  50                 return
  51         }
  52
  53         q.queued[key] = true
  54         q.merges = append(q.merges, key)
  55 }
  56
  57 // Next returns the next node that can start, and marks it as running.
  58 func (q *queue) Next() (nodeKey, bool) {
  59         if len(q.running) >= maxBuilds {
  60                 return nodeKey{}, false
  61         }
  62
  63         for _, id := range q.leaves[:min(len(q.leaves), leafWindow)] {
  64                 key := nodeKey{0, id}
  65                 if q.startable(key) {
  66                         q.running[key] = true
  67                         return key, true
  68                 }
  69         }
  70
  71         for _, key := range q.merges {
  72                 if q.startable(key) {
  73                         q.running[key] = true
  74                         return key, true
  75                 }
  76         }
  77
  78         return nodeKey{}, false
  79 }
  80
  81 func (q *queue) startable(key nodeKey) bool {
  82         return !q.running[key] && !q.failed[key]
  83 }
  84
  85 // Done removes a built node from the queue.
  86 func (q *queue) Done(key nodeKey) {
  87         delete(q.running, key)
  88         if key.level == 0 {
  89                 q.leaves = slices.DeleteFunc(q.leaves, func(id int) bool {
  90                         return id == key.index
  91                 })
  92
  93                 return
  94         }
  95
  96         delete(q.queued, key)
  97         q.merges = slices.DeleteFunc(q.merges, func(other nodeKey) bool {
  98                 return other == key
  99         })
 100 }
 101
 102 // Fail keeps a node in the queue, but it cannot start until Retry.
 103 func (q *queue) Fail(key nodeKey) {
 104         delete(q.running, key)
 105         q.failed[key] = true
 106 }
 107
 108 // Retry lets the failed nodes start again.
 109 func (q *queue) Retry() {
 110         clear(q.failed)
 111 }
 112
 113 // Running returns the number of builds that run.
 114 func (q *queue) Running() int {
 115         return len(q.running)
 116 }