X-Git-Url: http://git.giacomolli.uk/onechat/blobdiff_plain/e5343bd4c3ec69192e92f4be339b855eef36b6ec..c3c7bc90eb73e83f48fbeac6be149ade95a726c0:/memory.go?ds=sidebyside diff --git a/memory.go b/memory.go index eb28341..a240b6c 100644 --- a/memory.go +++ b/memory.go @@ -5,49 +5,46 @@ import ( "encoding/json" "errors" "fmt" + "log/slog" "os" "path/filepath" + "strings" "sync" "time" ) -// errNoView tells that a chat has messages but no saved view. -var errNoView = errors.New("the chat has messages but no saved view") +var ( + // errMemoryNoView tells that a chat has messages but no saved view. + errMemoryNoView = errors.New("the chat has messages but no saved view") -// Memory is the memory of one chat: the message log, the summary tree, the chat -// view, and the compaction view. The compaction view is the chat view merged -// further, at the limits of 16 to 32 KB. It takes the same new lines, and it -// merges again in one batch when it is too large or when the chat view merges. -// So the compactions share its prefix in the prompt cache. It saves the views -// to view.json after each change and loads them at start. It never builds the -// views again from the log, because a new view differs from the live one and so -// loses the prompt cache. -type Memory struct { - mutex sync.Mutex - root string - log *Log - tree *Tree - chat *View - cview *View -} - -// viewFile is the content of view.json. -type viewFile struct { - Chat viewState `json:"chat"` - Compaction viewState `json:"compaction"` -} + // errMemoryUnsettled tells that a line of the view stays pending, + // because its build failed. + errMemoryUnsettled = errors.New("a summary of the view failed, and it is tried again at the next message") +) -// openMemory opens the chat in root. The clock gives the dates of new messages -// and the days of new nodes. -func openMemory(ctx context.Context, root string, clock func() time.Time) (*Memory, error) { +// openMemory opens the chat in root and starts the builds that are ready. The +// clock gives the dates of new messages and the days of new nodes. logger gets +// one line for each failed build. The builds stop when ctx ends or at Close. +func openMemory(ctx context.Context, root string, clock func() time.Time, summarizer Summarizer, logger *slog.Logger) (*Memory, error) { log, err := openLog(ctx, root, clock) if err != nil { return nil, err } m := &Memory{ - root: root, - log: log, + root: root, + log: log, + queue: newBuildQueue(), + summarizer: summarizer, + logger: logger, + changed: make(chan struct{}), + } + + m.ctx, m.cancel = context.WithCancel(ctx) + + m.syncer, err = openSyncer(root) + if err != nil { + return nil, errors.Join(err, m.Close()) } m.tree, err = openTree(ctx, root, log.Len(), clock) @@ -60,41 +57,323 @@ func openMemory(ctx context.Context, root string, clock func() time.Time) (*Memo return nil, errors.Join(fmt.Errorf("load view: %w", err), m.Close()) } + err = m.restoreQueue() + if err != nil { + return nil, errors.Join(fmt.Errorf("restore build queue: %w", err), m.Close()) + } + + err = m.catchUp() + if err != nil { + return nil, errors.Join(err, m.Close()) + } + + m.mutex.Lock() + m.syncer.Written() + m.dispatch() + m.mutex.Unlock() + return m, nil } -// load reads the saved view. A crash can stop an append after the log write and -// before the view save, so load continues the view with the messages that it -// does not cover yet. +// Summarizer builds the text of one node with a model. +type Summarizer interface { + Summarize(ctx context.Context, task Task) (string, error) +} + +// Memory is the memory of one chat: the message log, the summary tree, the chat +// view, the compaction view, and the queue of the builds. The compaction view +// is the chat view merged further, at the limits of 16 to 32 KB. It takes the +// same new lines, and it merges again in one batch when it is too large or when +// the chat view merges. So the compactions share its prefix in the prompt +// cache. It saves the views to view.json after each change and loads them at +// start. It never builds the views again from the log, because a new view +// differs from the live one and so loses the prompt cache. +// +// Each write goes to the operating system at once, and a syncer makes the +// writes durable by group commit in the background. Flush waits until the +// writes so far are durable. +// +// The summarizer builds the nodes in the queue in the background. The start +// reads the indexes once to fill the queue again. After that, only a new +// message or a new node adds work. Memory keeps the nodes of the views and the +// tasks of the running builds. It reads all other messages and nodes from their +// day files through the indexes. +type Memory struct { + mutex sync.Mutex + root string + log *Log + syncer *Syncer + tree *Tree + chat *View + cview *View + queue *BuildQueue + summarizer Summarizer + logger *slog.Logger + ctx context.Context + cancel context.CancelFunc + workers sync.WaitGroup + changed chan struct{} +} + +// Close stops the builds, waits for them, makes the writes durable, closes the +// files of the chat, and releases its lock. +func (m *Memory) Close() error { + m.cancel() + m.workers.Wait() + + m.mutex.Lock() + defer m.mutex.Unlock() + + var err error + + if m.syncer != nil { + err = m.syncer.Close() + } + + if m.tree != nil { + err = errors.Join(err, m.tree.Close()) + } + + return errors.Join(err, m.log.Close()) +} + +// Append logs a new message, adds it to the tree and the views, and starts the +// builds that are ready, also the failed ones. +func (m *Memory) Append(kind, text string) (Message, error) { + return m.AppendDated(kind, text, "") +} + +// AppendDated appends a message with its RFC 3339 date and builds its nodes. An +// empty date uses the clock, as Append does. +func (m *Memory) AppendDated(kind, text, date string) (Message, error) { + m.mutex.Lock() + defer m.mutex.Unlock() + + return m.appendDated(kind, text, date) +} + +// appendDated appends under the memory lock and updates the views and builds. +func (m *Memory) appendDated(kind, text, date string) (Message, error) { + message, err := m.log.AppendDated(kind, text, date) + if err != nil { + return Message{}, fmt.Errorf("append dated log message: %w", err) + } + + defer m.syncer.Written() + + err = m.addMessage(message) + if err != nil { + return Message{}, fmt.Errorf("add dated message to memory: %w", err) + } + + err = m.save() + if err != nil { + return Message{}, fmt.Errorf("save dated message view: %w", err) + } + + m.queue.Retry() + m.dispatch() + + return message, nil +} + +// AppendOnce appends a queued message. prepare saves its log ID before the +// append, under the memory lock. A saved ID permits recovery after a stop +// between that save and the append, without a second log entry. prepare must +// not call Memory. The caller stops delivery after a failed append. +func (m *Memory) AppendOnce(kind, text string, id *int, prepare func(int) error) (Message, error) { + m.mutex.Lock() + defer m.mutex.Unlock() + + if id != nil && *id < m.log.Len() { + message, err := m.log.Message(*id) + if err != nil { + return Message{}, fmt.Errorf("read queued message: %w", err) + } + + if message.Kind != kind || message.Text != text { + return Message{}, errors.New("queued message does not match its log entry") + } + + return message, nil + } + + if id != nil && *id != m.log.Len() { + return Message{}, errors.New("queued message is beyond the log") + } + + err := prepare(m.log.Len()) + if err != nil { + return Message{}, fmt.Errorf("prepare queued message: %w", err) + } + + message, err := m.appendDated(kind, text, "") + if err != nil { + return Message{}, fmt.Errorf("append queued message: %w", err) + } + + return message, nil +} + +// children reads the two children of the node key. +func (m *Memory) children(key NodeKey) (Node, Node, error) { + var children [2]Node + + for n := range children { + child, built, err := m.tree.Node(key.level-1, 2*key.index+n) + if err != nil { + return Node{}, Node{}, err + } + + if !built { + return Node{}, Node{}, fmt.Errorf("node (%d, %d) is not built", key.level-1, 2*key.index+n) + } + + children[n] = child + } + + return children[0], children[1], nil +} + +// Wait returns when no build runs. The failed builds wait for the next message. +func (m *Memory) Wait(ctx context.Context) error { + for { + m.mutex.Lock() + + running := m.queue.Running() + changed := m.changed + m.mutex.Unlock() + + if running == 0 { + return nil + } + + select { + case <-changed: + case <-ctx.Done(): + return ctx.Err() + } + } +} + +// Settle returns when the chat view has no pending line, so that every message +// so far is summarized. A pending line with no build that runs belongs to a +// failed build, which waits for the next message, so Settle returns +// errMemoryUnsettled in place of a wait that never ends. +func (m *Memory) Settle(ctx context.Context) error { + for { + m.mutex.Lock() + + pending := m.chat.Pending() + running := m.queue.Running() + changed := m.changed + m.mutex.Unlock() + + if !pending { + return nil + } + + if running == 0 { + return errMemoryUnsettled + } + + select { + case <-changed: + case <-ctx.Done(): + return ctx.Err() + } + } +} + +// Flush returns when every write before the call is durable. +func (m *Memory) Flush(ctx context.Context) error { + return m.syncer.Flush(ctx) +} + +// Snapshot flushes all memory writes and runs snapshot while no append or +// summary write can change the log, indexes, tree, or views. snapshot must not +// call Memory. +func (m *Memory) Snapshot(ctx context.Context, snapshot func() error) error { + m.mutex.Lock() + defer m.mutex.Unlock() + + err := m.Flush(ctx) + if err != nil { + return fmt.Errorf("flush memory snapshot: %w", err) + } + + err = snapshot() + if err != nil { + return fmt.Errorf("copy memory snapshot: %w", err) + } + + return nil +} + +// ViewLines returns the lines of the chat view, each with its newline, without +// the leaf lines of the messages in skip. +func (m *Memory) ViewLines(skip map[int]bool) []string { + m.mutex.Lock() + defer m.mutex.Unlock() + + return m.chat.LinesWithout(skip) +} + +// Len returns the number of messages in the log. +func (m *Memory) Len() int { + m.mutex.Lock() + defer m.mutex.Unlock() + + return m.log.Len() +} + +// Message reads the message with the given ID. +func (m *Memory) Message(id int) (Message, error) { + m.mutex.Lock() + defer m.mutex.Unlock() + + return m.log.Message(id) +} + +// Node reads the node (l, i) if it is built. +func (m *Memory) Node(level, index int) (Node, bool, error) { + m.mutex.Lock() + defer m.mutex.Unlock() + + return m.tree.Node(level, index) +} + +// load reads the saved views. func (m *Memory) load() error { data, err := os.ReadFile(m.viewPath()) if errors.Is(err, os.ErrNotExist) && m.log.Len() == 0 { - m.chat = newView(chatHigh, chatLow) - m.cview = newView(compactionHigh, compactionLow) + m.chat = newView(viewChatHigh, viewChatLow) + m.cview = newView(viewCompactionHigh, viewCompactionLow) return nil } if errors.Is(err, os.ErrNotExist) { - return errNoView + return errMemoryNoView } if err != nil { return fmt.Errorf("read view: %w", err) } - var file viewFile + var file MemoryViewFile + err = json.Unmarshal(data, &file) if err != nil { return fmt.Errorf("decode view: %w", err) } - m.chat, err = restoreView(file.Chat, chatHigh, chatLow, m.tree) + m.chat, err = restoreView(file.Chat, viewChatHigh, viewChatLow, m.tree) if err != nil { return fmt.Errorf("chat view: %w", err) } - m.cview, err = restoreView(file.Compaction, compactionHigh, compactionLow, m.tree) + m.cview, err = restoreView(file.Compaction, viewCompactionHigh, viewCompactionLow, m.tree) if err != nil { return fmt.Errorf("compaction view: %w", err) } @@ -107,12 +386,73 @@ func (m *Memory) load() error { return fmt.Errorf("view covers %d messages, but the log has %d", m.chat.end, m.log.Len()) } + return nil +} + +// restoreQueue adds the work of the messages in the view to the queue. It +// builds each node that needs no model call and that a crash left out. +func (m *Memory) restoreQueue() error { + for id := range m.chat.end { + built, err := m.tree.Built(0, id) + if err != nil { + return err + } + + if built { + continue + } + + message, err := m.log.Message(id) + if err != nil { + return err + } + + leaf, free := freeLeaf(message) + if !free { + m.queue.AddLeaf(id) + + continue + } + + err = m.build(leaf) + if err != nil { + return err + } + } + + keys, err := m.tree.Keys() + if err != nil { + return err + } + + for _, key := range keys { + if key.index%2 != 0 { + continue + } + + err = m.settle(key) + if err != nil { + return err + } + } + + return nil +} + +// catchUp adds the messages that the view does not cover yet. A crash can stop +// an append after the log write and before the view save. +func (m *Memory) catchUp() error { if m.chat.end == m.log.Len() { return nil } for m.chat.end < m.log.Len() { - err = m.addMessage(m.log.Message(m.chat.end)) + message, err := m.log.Message(m.chat.end) + if err != nil { + return err + } + + err = m.addMessage(message) if err != nil { return err } @@ -121,111 +461,230 @@ func (m *Memory) load() error { return m.save() } -// Close closes the files of the chat and releases its lock. -func (m *Memory) Close() error { - m.mutex.Lock() - defer m.mutex.Unlock() +// dispatch starts the builds that are ready, up to the limit. A build whose +// input cannot be read fails and waits for the next message. +func (m *Memory) dispatch() { + for m.ctx.Err() == nil { + key, ready := m.queue.Next() + if !ready { + return + } - var err error - if m.tree != nil { - err = m.tree.Close() - } + task, err := m.task(key) + if err != nil { + m.failed(key, fmt.Errorf("read its input: %w", err)) + m.queue.Fail(key) - return errors.Join(err, m.log.Close()) -} + continue + } -// Append logs a new message and adds it to the tree and the view. -func (m *Memory) Append(kind, text string) (Message, error) { - m.mutex.Lock() - defer m.mutex.Unlock() + m.workers.Add(1) + + go m.run(key, task) + } +} - message, err := m.log.Append(kind, text) +// addMessage builds the leaf of message when it needs no model call, or else +// adds the leaf to the queue. Then it appends the leaf to the views. A crash can +// stop an append after the leaf build and before the view save, so a leaf that +// is built already stays as it is, and restoreQueue settles its parent. +func (m *Memory) addMessage(message Message) error { + built, err := m.tree.Built(0, message.ID) if err != nil { - return Message{}, err + return err } - err = m.addMessage(message) + leaf, free := freeLeaf(message) + switch { + case built: + case free: + err = m.build(leaf) + if err != nil { + return err + } + default: + m.queue.AddLeaf(message.ID) + } + + merged, err := m.chat.Append(m.tree) if err != nil { - return Message{}, err + return fmt.Errorf("append to chat view: %w", err) + } + + _, err = m.cview.Append(m.tree) + if err == nil && merged { + err = m.cview.Coarsen(m.chat, m.tree) } - err = m.save() if err != nil { - return Message{}, err + return fmt.Errorf("append to compaction view: %w", err) } - return message, nil + return nil } -// addMessage builds the leaf of message when it needs no model call and appends -// it to the view. -func (m *Memory) addMessage(message Message) error { - _, built := m.tree.Node(0, message.ID) - leaf, free := freeLeaf(message) - if free && !built { - err := m.build(leaf) - if err != nil { - return err - } +// save writes the views to view.json. +func (m *Memory) save() error { + data, err := json.Marshal(MemoryViewFile{ + Chat: m.chat.State(), + Compaction: m.cview.State(), + }) + if err != nil { + return fmt.Errorf("encode view: %w", err) } - merged := m.chat.Append(m.tree) - m.cview.Append(m.tree) - if merged { - m.cview.Coarsen(m.chat, m.tree) + err = writeFileAtomic(m.viewPath(), append(data, '\n')) + if err != nil { + return fmt.Errorf("save view: %w", err) } return nil } -// build adds node to the tree, and then each parent above it that needs no -// model call. +func (m *Memory) viewPath() string { + return filepath.Join(m.root, "view.json") +} + +// build adds node to the tree, keeps it in the views when it is a pending line, +// and settles its parent. func (m *Memory) build(node Node) error { + err := m.tree.Add(node) + if err != nil { + return err + } + + m.chat.Refresh(node) + m.cview.Refresh(node) + + return m.settle(NodeKey{node.Level, node.Index}) +} + +// settle looks at the parent of the built node key when both children are +// built. It builds each parent that needs no model call, up the tree, and adds +// the first other parent to the queue. +func (m *Memory) settle(key NodeKey) error { for { - err := m.tree.Add(node) + parent := NodeKey{key.level + 1, key.index / 2} + ready, err := m.ready(parent) + if err != nil || !ready { + return err + } + + left, right, err := m.children(parent) if err != nil { return err } - sibling, built := m.tree.Node(node.Level, node.Index^1) - if !built { + node, free := freeMerge(left, right) + if !free { + m.queue.AddMerge(parent) + return nil } - left, right := node, sibling - if node.Index%2 != 0 { - left, right = sibling, node + err = m.tree.Add(node) + if err != nil { + return err } - parent, free := freeMerge(left, right) - if !free { - return nil + key = parent + } +} + +// task returns the build task of the node key, with its input read from the log +// or the tree. +func (m *Memory) task(key NodeKey) (Task, error) { + task := Task{ + Level: key.level, + Index: key.index, + Context: m.context(key.level, key.index), + } + + if key.level == 0 { + message, err := m.log.Message(key.index) + if err != nil { + return Task{}, err } - node = parent + task.Input = []string{leafSource(message)} + + return task, nil } -} -// save writes the views to view.json. -func (m *Memory) save() error { - data, err := json.Marshal(viewFile{ - Chat: m.chat.State(), - Compaction: m.cview.State(), - }) + left, right, err := m.children(key) if err != nil { - return fmt.Errorf("encode view: %w", err) + return Task{}, err + } + + for _, child := range []Node{left, right} { + task.Input = append(task.Input, strings.ReplaceAll(child.Text, "\n", " ")) + } + + return task, nil +} + +// run builds one node with the summarizer. A failed build stays in the queue +// until the next message. +func (m *Memory) run(key NodeKey, task Task) { + defer m.workers.Done() + + text, err := m.summarizer.Summarize(m.ctx, task) + + m.mutex.Lock() + defer m.mutex.Unlock() + + if err == nil { + err = m.build(Node{ + Level: key.level, + Index: key.index, + Text: text, + Size: len(text), + }) + m.syncer.Written() + } + + if err != nil && m.ctx.Err() == nil { + m.failed(key, err) } - err = writeFileAtomic(m.viewPath(), append(data, '\n')) if err != nil { - return fmt.Errorf("save view: %w", err) + m.queue.Fail(key) + } else { + m.queue.Done(key) } - return nil + close(m.changed) + + m.changed = make(chan struct{}) + m.dispatch() } -func (m *Memory) viewPath() string { - return filepath.Join(m.root, "view.json") +// failed logs the failed build of the node key with its cause err. +func (m *Memory) failed(key NodeKey, err error) { + m.logger.Warn("summary failed", "level", key.level, "index", key.index, "err", err) +} + +// ready tells if both children of the node key are built and the node is not. +func (m *Memory) ready(key NodeKey) (bool, error) { + for _, index := range []int{2 * key.index, 2*key.index + 1} { + built, err := m.tree.Built(key.level-1, index) + if err != nil || !built { + return false, err + } + } + + built, err := m.tree.Built(key.level, key.index) + + return !built, err +} + +func (m *Memory) context(level, index int) string { + end := index + if level > 0 { + end = nodeEnd(level, index) + } + + return m.cview.Context(end) } // View returns the chat view as the model reads it. @@ -233,7 +692,7 @@ func (m *Memory) View() string { m.mutex.Lock() defer m.mutex.Unlock() - return m.chat.Render(m.tree) + return m.chat.Render() } // Context returns the compaction view for the build of node (l, i). It ends @@ -246,56 +705,92 @@ func (m *Memory) Context(level, index int) string { return m.context(level, index) } -func (m *Memory) context(level, index int) string { - end := index - if level > 0 { - end = nodeEnd(level, index) +// rebuildComplete checks the whole tree, not only the visible view. A compactor +// can fail without an HTTP error, for example with an empty line. Such a failed +// merge can be absent from the view and must prevent a swap. +func (m *Memory) rebuildComplete() error { + m.mutex.Lock() + defer m.mutex.Unlock() + + if len(m.queue.leaves) != 0 || len(m.queue.queued) != 0 { + return errors.New("rebuild failed: summary tree has unfinished nodes") } - return m.cview.Context(m.tree, end) + return nil } -// writeFileAtomic replaces the file at path with data. A reader or a crash sees -// the old file or the new file, never a part of one. -func writeFileAtomic(path string, data []byte) error { - dir := filepath.Dir(path) - file, err := os.CreateTemp(dir, "."+filepath.Base(path)+".*") +// buildModelLines returns the sizes of the nodes that the model wrote, in tree +// order. A node that fits its source with no model call is not one of them. +func (m *Memory) buildModelLines() ([]int, error) { + m.mutex.Lock() + defer m.mutex.Unlock() + + keys, err := m.tree.Keys() if err != nil { - return fmt.Errorf("create temporary file: %w", err) + return nil, fmt.Errorf("list built nodes: %w", err) } - _, err = file.Write(data) - if err == nil { - err = file.Sync() - } + var sizes []int - err = errors.Join(err, file.Close()) - if err == nil { - err = os.Rename(file.Name(), path) - } + for _, key := range keys { + var node Node - if err != nil { - return errors.Join(fmt.Errorf("replace file: %w", err), os.Remove(file.Name())) + node, _, err = m.tree.Node(key.level, key.index) + if err != nil { + return nil, fmt.Errorf("read built node: %w", err) + } + + var model bool + + model, err = m.buildBuiltByModel(key) + if err != nil { + return nil, fmt.Errorf("check model source: %w", err) + } + + if model { + sizes = append(sizes, node.Size) + } } - return syncDir(dir) + return sizes, nil } -// syncDir flushes the entries of dir to the disk. -func syncDir(dir string) (err error) { - handle, err := os.Open(dir) - if err != nil { - return fmt.Errorf("open directory: %w", err) - } +// buildBuiltByModel tells if the node key needed a model call: a leaf whose +// source is over nodeBytes, or a merge whose children do not fit in one node. +func (m *Memory) buildBuiltByModel(key NodeKey) (bool, error) { + if key.level == 0 { + message, err := m.log.Message(key.index) + if err != nil { + return false, fmt.Errorf("read leaf source: %w", err) + } - defer func() { - err = errors.Join(err, handle.Close()) - }() + _, free := freeLeaf(message) - err = handle.Sync() + return !free, nil + } + + left, right, err := m.children(key) if err != nil { - return fmt.Errorf("flush directory: %w", err) + return false, fmt.Errorf("read merge sources: %w", err) } - return nil + _, free := freeMerge(left, right) + + return !free, nil +} + +// Task is the build of the node (Level, Index). Context is the compaction view +// for the build. Input holds the source "kind: text" of a leaf, or the two +// child lines of a merge with their newlines as spaces. +type Task struct { + Level int + Index int + Context string + Input []string +} + +// MemoryViewFile is the content of view.json. +type MemoryViewFile struct { + Chat ViewState `json:"chat"` + Compaction ViewState `json:"compaction"` }