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

   1 package main
   2
   3 import (
   4         "cmp"
   5         "context"
   6         "encoding/binary"
   7         "encoding/json"
   8         "errors"
   9         "fmt"
  10         "math/bits"
  11         "os"
  12         "path/filepath"
  13         "slices"
  14         "time"
  15         "unicode/utf8"
  16 )
  17
  18 // nodeBytes is the size limit of a node text in UTF-8 bytes. A model can go a
  19 // few bytes over it, because the view measures the real sizes.
  20 const nodeBytes = 512
  21
  22 // treeRecordBytes is the size of a record of index/tree.idx: the position of
  23 // the node line, the node index, and the node level.
  24 const treeRecordBytes = 32
  25
  26 // Tree holds the built nodes of one chat. It writes one JSON line for each node
  27 // to tree/YYYY-MM-DD.jsonl, named for the local day of the build. Each node is
  28 // built once and never again. A node of a higher level is built only after both
  29 // of its children. After a failed write, the tree refuses all writes until it
  30 // is opened again. The caller holds the lock of the chat and serializes all
  31 // calls.
  32 //
  33 // The tree keeps no node in memory. The index index/tree.idx lists the node
  34 // lines in file order, and index/tree-L.idx holds the position of node (L, i)
  35 // in record i, with a zero record for a node that is not built. At open the
  36 // tree indexes the lines after the last indexed one, so a missing index is
  37 // built again from the day files.
  38 type Tree struct {
  39         clock   func() time.Time
  40         journal Journal
  41         dir     string
  42         order   *Records
  43         levels  map[int]*Records
  44         failure error
  45 }
  46
  47 // Node is one line of the summary tree. Node (l, i) covers the 2^l messages
  48 // from i*2^l on. A node of level 0 holds one message. A node of a higher level
  49 // merges its two children, (l-1, 2i) and (l-1, 2i+1). Size is the UTF-8 byte
  50 // length of Text.
  51 type Node struct {
  52         Level int    `json:"l"`
  53         Index int    `json:"i"`
  54         Text  string `json:"text"`
  55         Size  int    `json:"size"`
  56 }
  57
  58 // NodeKey names the node (l, i).
  59 type NodeKey struct {
  60         level int
  61         index int
  62 }
  63
  64 // openTree opens the tree of the chat in root and indexes its new lines. A chat
  65 // of total messages cannot have a node that covers a later message.
  66 func openTree(ctx context.Context, root string, total int, clock func() time.Time) (*Tree, error) {
  67         for _, dir := range []string{"tree", "index"} {
  68                 err := os.MkdirAll(filepath.Join(root, dir), 0o700)
  69                 if err != nil {
  70                         return nil, fmt.Errorf("create %s directory: %w", dir, err)
  71                 }
  72         }
  73
  74         t := &Tree{
  75                 clock: clock,
  76                 journal: Journal{
  77                         dir: filepath.Join(root, "tree"),
  78                 },
  79                 dir:    filepath.Join(root, "index"),
  80                 levels: make(map[int]*Records),
  81         }
  82
  83         err := t.open(ctx, total)
  84         if err != nil {
  85                 return nil, errors.Join(fmt.Errorf("index tree: %w", err), t.Close())
  86         }
  87
  88         return t, nil
  89 }
  90
  91 func (t *Tree) open(ctx context.Context, total int) error {
  92         var err error
  93
  94         t.order, err = openRecords(filepath.Join(t.dir, "tree.idx"), treeRecordBytes)
  95         if err != nil {
  96                 return fmt.Errorf("open node order index: %w", err)
  97         }
  98
  99         start := Position{}
 100         if t.order.count == 0 {
 101                 err = t.removeLevels()
 102                 if err != nil {
 103                         return fmt.Errorf("remove empty level indexes: %w", err)
 104                 }
 105         } else {
 106                 record, _, err := t.order.Read(t.order.count - 1)
 107                 if err != nil {
 108                         return fmt.Errorf("read last node order: %w", err)
 109                 }
 110
 111                 // A crash can stop an add between the two index writes, so the
 112                 // last node gets its level record again.
 113                 node, last := treeDecodeRecord(record)
 114                 err = t.writeLevel(node, last)
 115                 if err != nil {
 116                         return fmt.Errorf("restore last node level: %w", err)
 117                 }
 118
 119                 start = Position{
 120                         day:    last.day,
 121                         offset: last.offset + uint64(last.length),
 122                 }
 123         }
 124
 125         err = t.journal.Scan(ctx, start, func(line []byte, place Position) error {
 126                 var node Node
 127
 128                 err := json.Unmarshal(line, &node)
 129                 if err != nil {
 130                         return fmt.Errorf("decode tree node: %w", err)
 131                 }
 132
 133                 err = t.check(node)
 134                 if err != nil {
 135                         return fmt.Errorf("validate scanned node: %w", err)
 136                 }
 137
 138                 if nodeEnd(node.Level, node.Index) > total {
 139                         return fmt.Errorf("node (%d, %d) covers messages after the log", node.Level, node.Index)
 140                 }
 141
 142                 err = t.addIndex(NodeKey{node.Level, node.Index}, place)
 143                 if err != nil {
 144                         return fmt.Errorf("index scanned node: %w", err)
 145                 }
 146
 147                 return nil
 148         })
 149         if err != nil {
 150                 return fmt.Errorf("scan node journal: %w", err)
 151         }
 152
 153         return nil
 154 }
 155
 156 // Close closes the files of the tree.
 157 func (t *Tree) Close() error {
 158         err := t.journal.Close()
 159         if t.order != nil {
 160                 err = errors.Join(err, t.order.Close())
 161                 t.order = nil
 162         }
 163
 164         for level, index := range t.levels {
 165                 err = errors.Join(err, index.Close())
 166
 167                 delete(t.levels, level)
 168         }
 169
 170         return err
 171 }
 172
 173 // removeLevels removes the level indexes, which are not valid without the order
 174 // index.
 175 func (t *Tree) removeLevels() error {
 176         paths, err := filepath.Glob(filepath.Join(t.dir, "tree-*.idx"))
 177         if err != nil {
 178                 return fmt.Errorf("list level indexes: %w", err)
 179         }
 180
 181         for _, path := range paths {
 182                 err = os.Remove(path)
 183                 if err != nil {
 184                         return fmt.Errorf("remove level index: %w", err)
 185                 }
 186         }
 187
 188         return nil
 189 }
 190
 191 // writeLevel writes the position of node to its level index.
 192 func (t *Tree) writeLevel(node NodeKey, place Position) error {
 193         level, err := t.level(node.level)
 194         if err != nil {
 195                 return fmt.Errorf("open node level index: %w", err)
 196         }
 197
 198         record := make([]byte, positionBytes)
 199         encodePosition(record, place)
 200
 201         err = level.Write(node.index, record)
 202         if err != nil {
 203                 return fmt.Errorf("write node level record: %w", err)
 204         }
 205
 206         return nil
 207 }
 208
 209 // check makes sure that node is valid and can join the tree now.
 210 func (t *Tree) check(node Node) error {
 211         if node.Level < 0 || node.Level >= bits.UintSize-2 || node.Index < 0 {
 212                 return fmt.Errorf("node (%d, %d) is out of range", node.Level, node.Index)
 213         }
 214
 215         if node.Text == "" || !utf8.ValidString(node.Text) || node.Size != len(node.Text) {
 216                 return fmt.Errorf("node (%d, %d) has no valid text of its size", node.Level, node.Index)
 217         }
 218
 219         built, err := t.Built(node.Level, node.Index)
 220         if err != nil {
 221                 return fmt.Errorf("check target node: %w", err)
 222         }
 223
 224         if built {
 225                 return fmt.Errorf("node (%d, %d) is already built", node.Level, node.Index)
 226         }
 227
 228         if node.Level == 0 {
 229                 return nil
 230         }
 231
 232         for _, index := range []int{2 * node.Index, 2*node.Index + 1} {
 233                 built, err = t.Built(node.Level-1, index)
 234                 if err != nil {
 235                         return fmt.Errorf("check child node: %w", err)
 236                 }
 237
 238                 if !built {
 239                         return fmt.Errorf("node (%d, %d) has a child that is not built", node.Level, node.Index)
 240                 }
 241         }
 242
 243         return nil
 244 }
 245
 246 // addIndex adds the node line at place to both indexes.
 247 func (t *Tree) addIndex(node NodeKey, place Position) error {
 248         record := make([]byte, treeRecordBytes)
 249         encodePosition(record, place)
 250         binary.LittleEndian.PutUint64(record[positionBytes:], uint64(node.index))
 251         binary.LittleEndian.PutUint32(record[positionBytes+8:], uint32(node.level))
 252
 253         err := t.order.Append(record)
 254         if err != nil {
 255                 return fmt.Errorf("append node order: %w", err)
 256         }
 257
 258         err = t.writeLevel(node, place)
 259         if err != nil {
 260                 return fmt.Errorf("write indexed node level: %w", err)
 261         }
 262
 263         return nil
 264 }
 265
 266 // level returns the index of level, and opens it at first use.
 267 func (t *Tree) level(level int) (*Records, error) {
 268         index, open := t.levels[level]
 269         if open {
 270                 return index, nil
 271         }
 272
 273         index, err := openRecords(filepath.Join(t.dir, fmt.Sprintf("tree-%d.idx", level)), positionBytes)
 274         if err != nil {
 275                 return nil, err
 276         }
 277
 278         t.levels[level] = index
 279
 280         return index, nil
 281 }
 282
 283 // Built tells if the node (l, i) is built.
 284 func (t *Tree) Built(level, index int) (bool, error) {
 285         _, built, err := t.place(level, index)
 286         return built, err
 287 }
 288
 289 // place returns the position of the node (l, i) if it is built.
 290 func (t *Tree) place(level, index int) (Position, bool, error) {
 291         if level < 0 || level >= bits.UintSize-2 || index < 0 {
 292                 return Position{}, false, nil
 293         }
 294
 295         records, err := t.level(level)
 296         if err != nil {
 297                 return Position{}, false, fmt.Errorf("open node level index: %w", err)
 298         }
 299
 300         record, found, err := records.Read(index)
 301         if err != nil || !found {
 302                 if err != nil {
 303                         return Position{}, false, fmt.Errorf("read node level record: %w", err)
 304                 }
 305
 306                 return Position{}, false, nil
 307         }
 308
 309         place := decodePosition(record)
 310
 311         return place, place.length > 0, nil
 312 }
 313
 314 // Add writes a new node to the tree.
 315 func (t *Tree) Add(node Node) error {
 316         if t.failure != nil {
 317                 return fmt.Errorf("tree needs a reopen after a failed write: %w", t.failure)
 318         }
 319
 320         err := t.check(node)
 321         if err != nil {
 322                 return fmt.Errorf("validate node: %w", err)
 323         }
 324
 325         line, err := json.Marshal(node)
 326         if err != nil {
 327                 return fmt.Errorf("encode node: %w", err)
 328         }
 329
 330         place, err := t.journal.Write(t.clock().Format(time.DateOnly), append(line, '\n'))
 331         if err == nil {
 332                 err = t.addIndex(NodeKey{node.Level, node.Index}, place)
 333         }
 334
 335         if err != nil {
 336                 t.failure = err
 337                 return fmt.Errorf("write node (%d, %d): %w", node.Level, node.Index, err)
 338         }
 339
 340         return nil
 341 }
 342
 343 // Node reads the node (l, i) if it is built.
 344 func (t *Tree) Node(level, index int) (Node, bool, error) {
 345         place, built, err := t.place(level, index)
 346         if err != nil || !built {
 347                 if err != nil {
 348                         return Node{}, false, fmt.Errorf("find node position: %w", err)
 349                 }
 350
 351                 return Node{}, false, nil
 352         }
 353
 354         line, err := t.journal.Read(place)
 355         if err != nil {
 356                 return Node{}, false, fmt.Errorf("read node (%d, %d): %w", level, index, err)
 357         }
 358
 359         var node Node
 360
 361         err = json.Unmarshal(line, &node)
 362         if err != nil {
 363                 return Node{}, false, fmt.Errorf("decode node (%d, %d): %w", level, index, err)
 364         }
 365
 366         if node.Level != level || node.Index != index {
 367                 return Node{}, false, fmt.Errorf("index of node (%d, %d) points to node (%d, %d)", level, index, node.Level, node.Index)
 368         }
 369
 370         return node, true, nil
 371 }
 372
 373 // Keys returns the keys of the built nodes, ordered by level and then by index.
 374 func (t *Tree) Keys() ([]NodeKey, error) {
 375         keys := make([]NodeKey, 0, t.order.count)
 376         for n := range t.order.count {
 377                 record, _, err := t.order.Read(n)
 378                 if err != nil {
 379                         return nil, err
 380                 }
 381
 382                 key, _ := treeDecodeRecord(record)
 383                 keys = append(keys, key)
 384         }
 385
 386         slices.SortFunc(keys, func(a, b NodeKey) int {
 387                 return cmp.Or(cmp.Compare(a.level, b.level), cmp.Compare(a.index, b.index))
 388         })
 389
 390         return keys, nil
 391 }
 392
 393 // Count returns the number of built nodes.
 394 func (t *Tree) Count() int {
 395         return t.order.count
 396 }
 397
 398 // treeDecodeRecord reads the key and the position of a record of
 399 // index/tree.idx.
 400 func treeDecodeRecord(record []byte) (NodeKey, Position) {
 401         key := NodeKey{
 402                 level: int(binary.LittleEndian.Uint32(record[positionBytes+8:])),
 403                 index: int(binary.LittleEndian.Uint64(record[positionBytes:])),
 404         }
 405
 406         return key, decodePosition(record)
 407 }
 408
 409 // nodeEnd returns the ID after the last message of node (l, i).
 410 func nodeEnd(level, index int) int {
 411         return (index + 1) << level
 412 }
 413
 414 // freeLeaf returns a text-only leaf when its source fits in a node. Images need
 415 // a model description even when their text fits.
 416 func freeLeaf(message Message) (Node, bool) {
 417         text := leafSource(message)
 418         if len(text) > nodeBytes || len(imageFiles(message.Text)) > 0 {
 419                 return Node{}, false
 420         }
 421
 422         return Node{
 423                 Level: 0,
 424                 Index: message.ID,
 425                 Text:  text,
 426                 Size:  len(text),
 427         }, true
 428 }
 429
 430 // leafSource returns the text that the leaf of message holds or summarizes.
 431 func leafSource(message Message) string {
 432         return message.Kind + ": " + message.Text
 433 }
 434
 435 // freeMerge returns the parent of two sibling nodes when both texts, joined by
 436 // a newline, fit in a node. Such a parent needs no model call.
 437 func freeMerge(left, right Node) (Node, bool) {
 438         if left.Size+1+right.Size > nodeBytes {
 439                 return Node{}, false
 440         }
 441
 442         text := left.Text + "\n" + right.Text
 443
 444         return Node{
 445                 Level: left.Level + 1,
 446                 Index: left.Index / 2,
 447                 Text:  text,
 448                 Size:  len(text),
 449         }, true
 450 }