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