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

   1 package main
   2
   3 import (
   4         "context"
   5         "encoding/json"
   6         "errors"
   7         "fmt"
   8         "os"
   9         "path/filepath"
  10         "sync"
  11         "time"
  12 )
  13
  14 // errNoView tells that a chat has messages but no saved view.
  15 var errNoView = errors.New("the chat has messages but no saved view")
  16
  17 // Memory is the memory of one chat: the message log, the summary tree, the chat
  18 // view, and the compaction view. The compaction view is the chat view merged
  19 // further, at the limits of 16 to 32 KB. It takes the same new lines, and it
  20 // merges again in one batch when it is too large or when the chat view merges.
  21 // So the compactions share its prefix in the prompt cache. It saves the views
  22 // to view.json after each change and loads them at start. It never builds the
  23 // views again from the log, because a new view differs from the live one and so
  24 // loses the prompt cache.
  25 type Memory struct {
  26         mutex sync.Mutex
  27         root  string
  28         log   *Log
  29         tree  *Tree
  30         chat  *View
  31         cview *View
  32 }
  33
  34 // viewFile is the content of view.json.
  35 type viewFile struct {
  36         Chat       viewState `json:"chat"`
  37         Compaction viewState `json:"compaction"`
  38 }
  39
  40 // openMemory opens the chat in root. The clock gives the dates of new messages
  41 // and the days of new nodes.
  42 func openMemory(ctx context.Context, root string, clock func() time.Time) (*Memory, error) {
  43         log, err := openLog(ctx, root, clock)
  44         if err != nil {
  45                 return nil, err
  46         }
  47
  48         m := &Memory{
  49                 root: root,
  50                 log:  log,
  51         }
  52
  53         m.tree, err = openTree(ctx, root, log.Len(), clock)
  54         if err != nil {
  55                 return nil, errors.Join(err, m.Close())
  56         }
  57
  58         err = m.load()
  59         if err != nil {
  60                 return nil, errors.Join(fmt.Errorf("load view: %w", err), m.Close())
  61         }
  62
  63         return m, nil
  64 }
  65
  66 // load reads the saved view. A crash can stop an append after the log write and
  67 // before the view save, so load continues the view with the messages that it
  68 // does not cover yet.
  69 func (m *Memory) load() error {
  70         data, err := os.ReadFile(m.viewPath())
  71         if errors.Is(err, os.ErrNotExist) && m.log.Len() == 0 {
  72                 m.chat = newView(chatHigh, chatLow)
  73                 m.cview = newView(compactionHigh, compactionLow)
  74
  75                 return nil
  76         }
  77
  78         if errors.Is(err, os.ErrNotExist) {
  79                 return errNoView
  80         }
  81
  82         if err != nil {
  83                 return fmt.Errorf("read view: %w", err)
  84         }
  85
  86         var file viewFile
  87         err = json.Unmarshal(data, &file)
  88         if err != nil {
  89                 return fmt.Errorf("decode view: %w", err)
  90         }
  91
  92         m.chat, err = restoreView(file.Chat, chatHigh, chatLow, m.tree)
  93         if err != nil {
  94                 return fmt.Errorf("chat view: %w", err)
  95         }
  96
  97         m.cview, err = restoreView(file.Compaction, compactionHigh, compactionLow, m.tree)
  98         if err != nil {
  99                 return fmt.Errorf("compaction view: %w", err)
 100         }
 101
 102         if m.cview.end != m.chat.end {
 103                 return fmt.Errorf("compaction view covers %d messages, but the chat view %d", m.cview.end, m.chat.end)
 104         }
 105
 106         if m.chat.end > m.log.Len() {
 107                 return fmt.Errorf("view covers %d messages, but the log has %d", m.chat.end, m.log.Len())
 108         }
 109
 110         if m.chat.end == m.log.Len() {
 111                 return nil
 112         }
 113
 114         for m.chat.end < m.log.Len() {
 115                 err = m.addMessage(m.log.Message(m.chat.end))
 116                 if err != nil {
 117                         return err
 118                 }
 119         }
 120
 121         return m.save()
 122 }
 123
 124 // Close closes the files of the chat and releases its lock.
 125 func (m *Memory) Close() error {
 126         m.mutex.Lock()
 127         defer m.mutex.Unlock()
 128
 129         var err error
 130         if m.tree != nil {
 131                 err = m.tree.Close()
 132         }
 133
 134         return errors.Join(err, m.log.Close())
 135 }
 136
 137 // Append logs a new message and adds it to the tree and the view.
 138 func (m *Memory) Append(kind, text string) (Message, error) {
 139         m.mutex.Lock()
 140         defer m.mutex.Unlock()
 141
 142         message, err := m.log.Append(kind, text)
 143         if err != nil {
 144                 return Message{}, err
 145         }
 146
 147         err = m.addMessage(message)
 148         if err != nil {
 149                 return Message{}, err
 150         }
 151
 152         err = m.save()
 153         if err != nil {
 154                 return Message{}, err
 155         }
 156
 157         return message, nil
 158 }
 159
 160 // addMessage builds the leaf of message when it needs no model call and appends
 161 // it to the view.
 162 func (m *Memory) addMessage(message Message) error {
 163         _, built := m.tree.Node(0, message.ID)
 164         leaf, free := freeLeaf(message)
 165         if free && !built {
 166                 err := m.build(leaf)
 167                 if err != nil {
 168                         return err
 169                 }
 170         }
 171
 172         merged := m.chat.Append(m.tree)
 173         m.cview.Append(m.tree)
 174         if merged {
 175                 m.cview.Coarsen(m.chat, m.tree)
 176         }
 177
 178         return nil
 179 }
 180
 181 // build adds node to the tree, and then each parent above it that needs no
 182 // model call.
 183 func (m *Memory) build(node Node) error {
 184         for {
 185                 err := m.tree.Add(node)
 186                 if err != nil {
 187                         return err
 188                 }
 189
 190                 sibling, built := m.tree.Node(node.Level, node.Index^1)
 191                 if !built {
 192                         return nil
 193                 }
 194
 195                 left, right := node, sibling
 196                 if node.Index%2 != 0 {
 197                         left, right = sibling, node
 198                 }
 199
 200                 parent, free := freeMerge(left, right)
 201                 if !free {
 202                         return nil
 203                 }
 204
 205                 node = parent
 206         }
 207 }
 208
 209 // save writes the views to view.json.
 210 func (m *Memory) save() error {
 211         data, err := json.Marshal(viewFile{
 212                 Chat:       m.chat.State(),
 213                 Compaction: m.cview.State(),
 214         })
 215         if err != nil {
 216                 return fmt.Errorf("encode view: %w", err)
 217         }
 218
 219         err = writeFileAtomic(m.viewPath(), append(data, '\n'))
 220         if err != nil {
 221                 return fmt.Errorf("save view: %w", err)
 222         }
 223
 224         return nil
 225 }
 226
 227 func (m *Memory) viewPath() string {
 228         return filepath.Join(m.root, "view.json")
 229 }
 230
 231 // View returns the chat view as the model reads it.
 232 func (m *Memory) View() string {
 233         m.mutex.Lock()
 234         defer m.mutex.Unlock()
 235
 236         return m.chat.Render(m.tree)
 237 }
 238
 239 // Context returns the compaction view for the build of node (l, i). It ends
 240 // before the message of a leaf, or after the last message of a merge, and it
 241 // stops at the first line that is not built.
 242 func (m *Memory) Context(level, index int) string {
 243         m.mutex.Lock()
 244         defer m.mutex.Unlock()
 245
 246         return m.context(level, index)
 247 }
 248
 249 func (m *Memory) context(level, index int) string {
 250         end := index
 251         if level > 0 {
 252                 end = nodeEnd(level, index)
 253         }
 254
 255         return m.cview.Context(m.tree, end)
 256 }
 257
 258 // writeFileAtomic replaces the file at path with data. A reader or a crash sees
 259 // the old file or the new file, never a part of one.
 260 func writeFileAtomic(path string, data []byte) error {
 261         dir := filepath.Dir(path)
 262         file, err := os.CreateTemp(dir, "."+filepath.Base(path)+".*")
 263         if err != nil {
 264                 return fmt.Errorf("create temporary file: %w", err)
 265         }
 266
 267         _, err = file.Write(data)
 268         if err == nil {
 269                 err = file.Sync()
 270         }
 271
 272         err = errors.Join(err, file.Close())
 273         if err == nil {
 274                 err = os.Rename(file.Name(), path)
 275         }
 276
 277         if err != nil {
 278                 return errors.Join(fmt.Errorf("replace file: %w", err), os.Remove(file.Name()))
 279         }
 280
 281         return syncDir(dir)
 282 }
 283
 284 // syncDir flushes the entries of dir to the disk.
 285 func syncDir(dir string) (err error) {
 286         handle, err := os.Open(dir)
 287         if err != nil {
 288                 return fmt.Errorf("open directory: %w", err)
 289         }
 290
 291         defer func() {
 292                 err = errors.Join(err, handle.Close())
 293         }()
 294
 295         err = handle.Sync()
 296         if err != nil {
 297                 return fmt.Errorf("flush directory: %w", err)
 298         }
 299
 300         return nil
 301 }