From: Anderson A. Giacomolli Date: Fri, 6 Nov 2026 14:44:12 +0000 (-0300) Subject: Build the nodes from a queue in the background X-Git-Url: http://git.giacomolli.uk/onechat/commitdiff_plain/e754189bc2aa32ec43a5dc4ee01d7a6c6e4e388c?ds=sidebyside;hp=e5343bd4c3ec69192e92f4be339b855eef36b6ec Build the nodes from a queue in the background The core keeps the nodes that are ready for a model build in a queue and never scans the tree for work. A leaf joins the queue when its message is too long to be its own node. A merge joins when both of its children are built and their texts do not fit in one node. A parent that fits is built at once, up the tree. At most 8 builds run at once. A leaf starts once fewer than 8 leaves before it are unbuilt. Leaves start before merges, oldest first, because a turn waits for the leaves of all earlier messages. A failed build starts again at the next message. Each build gets its compaction view and its input, the source of a leaf or the two child lines of a merge. The model stays behind a small Summarizer interface, and the tests use a fake one. The start reads the log and the tree once to fill the queue again, and builds the free nodes that a crash left out. Issue: beb9a70 --- diff --git a/memory.go b/memory.go index eb28341..9c8073f 100644 --- a/memory.go +++ b/memory.go @@ -7,6 +7,7 @@ import ( "fmt" "os" "path/filepath" + "strings" "sync" "time" ) @@ -14,21 +15,46 @@ import ( // errNoView tells that a chat has messages but no saved view. var errNoView = errors.New("the chat has messages but no saved view") +// Summarizer builds the text of one node with a model. +type Summarizer interface { + Summarize(ctx context.Context, task Task) (string, error) +} + +// 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 +} + // 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. +// 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. +// +// The summarizer builds the nodes in the queue in the background. The start +// reads the log and the tree once to fill the queue again. After that, only a +// new message or a new node adds work. type Memory struct { - mutex sync.Mutex - root string - log *Log - tree *Tree - chat *View - cview *View + mutex sync.Mutex + root string + log *Log + tree *Tree + chat *View + cview *View + queue *queue + summarizer Summarizer + ctx context.Context + cancel context.CancelFunc + workers sync.WaitGroup + changed chan struct{} } // viewFile is the content of view.json. @@ -37,19 +63,25 @@ type viewFile struct { Compaction viewState `json:"compaction"` } -// 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. The builds +// stop when ctx ends or at Close. +func openMemory(ctx context.Context, root string, clock func() time.Time, summarizer Summarizer) (*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: newQueue(), + summarizer: summarizer, + changed: make(chan struct{}), } + m.ctx, m.cancel = context.WithCancel(ctx) + m.tree, err = openTree(ctx, root, log.Len(), clock) if err != nil { return nil, errors.Join(err, m.Close()) @@ -60,12 +92,24 @@ 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.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. +// 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 { @@ -107,12 +151,52 @@ 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 { + if m.tree.Built(0, id) { + continue + } + + leaf, free := freeLeaf(m.log.Message(id)) + if !free { + m.queue.AddLeaf(id) + continue + } + + err := m.build(leaf) + if err != nil { + return err + } + } + + for _, node := range m.tree.Nodes() { + if node.Index%2 != 0 { + continue + } + + err := m.settle(node) + 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)) + err := m.addMessage(m.log.Message(m.chat.end)) if err != nil { return err } @@ -121,8 +205,12 @@ func (m *Memory) load() error { return m.save() } -// Close closes the files of the chat and releases its lock. +// Close stops the builds, waits for them, 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() @@ -134,7 +222,8 @@ func (m *Memory) Close() error { return errors.Join(err, m.log.Close()) } -// Append logs a new message and adds it to the tree and the view. +// 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) { m.mutex.Lock() defer m.mutex.Unlock() @@ -154,19 +243,23 @@ func (m *Memory) Append(kind, text string) (Message, error) { return Message{}, err } + m.queue.Retry() + m.dispatch() + return message, nil } -// addMessage builds the leaf of message when it needs no model call and appends -// it to the view. +// 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. func (m *Memory) addMessage(message Message) error { - _, built := m.tree.Node(0, message.ID) leaf, free := freeLeaf(message) - if free && !built { + if free { err := m.build(leaf) if err != nil { return err } + } else { + m.queue.AddLeaf(message.ID) } merged := m.chat.Append(m.tree) @@ -178,17 +271,23 @@ func (m *Memory) addMessage(message Message) error { return nil } -// build adds node to the tree, and then each parent above it that needs no -// model call. +// build adds node to the tree and settles its parent. func (m *Memory) build(node Node) error { - for { - err := m.tree.Add(node) - if err != nil { - return err - } + err := m.tree.Add(node) + if err != nil { + return err + } + + return m.settle(node) +} +// settle looks at the parent of node 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(node Node) error { + for { sibling, built := m.tree.Node(node.Level, node.Index^1) - if !built { + if !built || m.tree.Built(node.Level+1, node.Index/2) { return nil } @@ -199,13 +298,104 @@ func (m *Memory) build(node Node) error { parent, free := freeMerge(left, right) if !free { + m.queue.AddMerge(nodeKey{node.Level + 1, node.Index / 2}) return nil } + err := m.tree.Add(parent) + if err != nil { + return err + } + node = parent } } +// dispatch starts the builds that are ready, up to the limit. +func (m *Memory) dispatch() { + for m.ctx.Err() == nil { + key, ready := m.queue.Next() + if !ready { + return + } + + task := m.task(key) + m.workers.Add(1) + go m.run(key, task) + } +} + +// task returns the build task of the node key. +func (m *Memory) task(key nodeKey) Task { + task := Task{ + Level: key.level, + Index: key.index, + Context: m.context(key.level, key.index), + } + + if key.level == 0 { + task.Input = []string{leafSource(m.log.Message(key.index))} + return task + } + + for _, index := range []int{2 * key.index, 2*key.index + 1} { + child, _ := m.tree.Node(key.level-1, index) + task.Input = append(task.Input, strings.ReplaceAll(child.Text, "\n", " ")) + } + + return task +} + +// 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), + }) + } + + if err != nil { + m.queue.Fail(key) + } else { + m.queue.Done(key) + } + + close(m.changed) + m.changed = make(chan struct{}) + m.dispatch() +} + +// 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() + } + } +} + // save writes the views to view.json. func (m *Memory) save() error { data, err := json.Marshal(viewFile{ diff --git a/memory_test.go b/memory_test.go index 426654a..e868220 100644 --- a/memory_test.go +++ b/memory_test.go @@ -1,10 +1,17 @@ package main import ( + "context" + "errors" "fmt" + "maps" "os" "path/filepath" + "slices" + "strings" + "sync" "testing" + "unicode/utf8" "github.com/stretchr/testify/require" ) @@ -12,7 +19,7 @@ import ( func TestMemoryRestartLoadsTheSameView(t *testing.T) { root := t.TempDir() clock := newTestClock() - memory, err := openMemory(t.Context(), root, clock.Now) + memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer()) require.NoError(t, err) memory.chat.high = 200 @@ -31,7 +38,7 @@ func TestMemoryRestartLoadsTheSameView(t *testing.T) { require.Less(t, len(memory.chat.lines), 100) require.NoError(t, memory.Close()) - memory, err = openMemory(t.Context(), root, clock.Now) + memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer()) require.NoError(t, err) defer func() { @@ -47,7 +54,7 @@ func TestMemoryRestartLoadsTheSameView(t *testing.T) { } func TestMemoryKeepsTheCompactionViewCoarser(t *testing.T) { - memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now) + memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, newFakeSummarizer()) require.NoError(t, err) defer func() { @@ -75,8 +82,125 @@ func TestMemoryKeepsTheCompactionViewCoarser(t *testing.T) { } } +func TestMemoryRunsAtMostEightBuilds(t *testing.T) { + summarizer := newFakeSummarizer() + summarizer.gate = make(chan struct{}) + memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer) + require.NoError(t, err) + + defer func() { + require.NoError(t, memory.Close()) + }() + + for n := range 12 { + _, err = memory.Append("echo", longText(n)) + require.NoError(t, err) + } + + memory.mutex.Lock() + running := slices.Collect(maps.Keys(memory.queue.running)) + memory.mutex.Unlock() + slices.SortFunc(running, func(a, b nodeKey) int { + return a.index - b.index + }) + require.Equal(t, []nodeKey{{0, 0}, {0, 1}, {0, 2}, {0, 3}, {0, 4}, {0, 5}, {0, 6}, {0, 7}}, running) + + close(summarizer.gate) + require.NoError(t, memory.Wait(t.Context())) + require.Equal(t, maxBuilds, summarizer.peak) + + for level, count := range []int{12, 6, 3, 1} { + for index := range count { + require.True(t, memory.tree.Built(level, index), "node (%d, %d)", level, index) + } + } +} + +func TestMemoryStartsALeafAfterFewerThanEightUnbuilt(t *testing.T) { + summarizer := newFakeSummarizer() + summarizer.SetFailing(true) + memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer) + require.NoError(t, err) + + defer func() { + require.NoError(t, memory.Close()) + }() + + for n := range 9 { + _, err = memory.Append("echo", longText(n)) + require.NoError(t, err) + } + + require.NoError(t, memory.Wait(t.Context())) + for _, task := range summarizer.Tasks() { + require.Less(t, task.Index, 8) + } + + require.False(t, memory.tree.Built(0, 8)) + summarizer.SetFailing(false) + + _, err = memory.Append("user", "retry") + require.NoError(t, err) + require.NoError(t, memory.Wait(t.Context())) + require.True(t, memory.tree.Built(0, 8)) + require.True(t, memory.tree.Built(3, 0)) +} + +func TestMemoryMergesAfterBothHalves(t *testing.T) { + summarizer := newFakeSummarizer() + memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer) + require.NoError(t, err) + + defer func() { + require.NoError(t, memory.Close()) + }() + + _, err = memory.Append("echo", longText(0)) + require.NoError(t, err) + require.NoError(t, memory.Wait(t.Context())) + + _, err = memory.Append("echo", longText(1)) + require.NoError(t, err) + require.NoError(t, memory.Wait(t.Context())) + require.Len(t, summarizer.Tasks(), 3) + + left, _ := memory.tree.Node(0, 0) + right, _ := memory.tree.Node(0, 1) + merge := summarizer.Tasks()[2] + require.Equal(t, []string{left.Text, right.Text}, merge.Input) + require.Equal(t, "\n0+1|"+left.Text+"\n1+1|"+right.Text+"\n", merge.Context) + require.True(t, memory.tree.Built(1, 0)) +} + +func TestMemoryRestartRestoresTheQueue(t *testing.T) { + root := t.TempDir() + clock := newTestClock() + summarizer := newFakeSummarizer() + summarizer.SetFailing(true) + memory, err := openMemory(t.Context(), root, clock.Now, summarizer) + require.NoError(t, err) + + for n := range 4 { + _, err = memory.Append("echo", longText(n)) + require.NoError(t, err) + } + + require.NoError(t, memory.Wait(t.Context())) + require.NoError(t, memory.Close()) + + memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer()) + require.NoError(t, err) + + defer func() { + require.NoError(t, memory.Close()) + }() + + require.NoError(t, memory.Wait(t.Context())) + require.True(t, memory.tree.Built(2, 0)) +} + func TestMemoryBuildsFreeParents(t *testing.T) { - memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now) + memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, newFakeSummarizer()) require.NoError(t, err) defer func() { @@ -94,23 +218,32 @@ func TestMemoryBuildsFreeParents(t *testing.T) { } func TestMemoryKeepsALongMessageOutOfTheView(t *testing.T) { - memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now) + summarizer := newFakeSummarizer() + summarizer.SetFailing(true) + memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer) require.NoError(t, err) defer func() { require.NoError(t, memory.Close()) }() - _, err = memory.Append("echo", string(make([]byte, nodeBytes))) + _, err = memory.Append("echo", longText(0)) require.NoError(t, err) - + require.NoError(t, memory.Wait(t.Context())) require.Equal(t, "\n0+1|"+pending+"\n", memory.View()) + + summarizer.SetFailing(false) + + _, err = memory.Append("user", "again") + require.NoError(t, err) + require.NoError(t, memory.Wait(t.Context())) + require.Equal(t, "\n0+1|"+fakeLine(memory.task(nodeKey{0, 0}))+"\n1+1|user: again\n", memory.View()) } func TestMemoryCoversMessagesAfterTheSavedView(t *testing.T) { root := t.TempDir() clock := newTestClock() - memory, err := openMemory(t.Context(), root, clock.Now) + memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer()) require.NoError(t, err) _, err = memory.Append("user", "a") @@ -120,7 +253,7 @@ func TestMemoryCoversMessagesAfterTheSavedView(t *testing.T) { require.NoError(t, err) require.NoError(t, memory.Close()) - memory, err = openMemory(t.Context(), root, clock.Now) + memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer()) require.NoError(t, err) defer func() { @@ -133,7 +266,7 @@ func TestMemoryCoversMessagesAfterTheSavedView(t *testing.T) { func TestMemoryRefusesToRebuildAMissingView(t *testing.T) { root := t.TempDir() clock := newTestClock() - memory, err := openMemory(t.Context(), root, clock.Now) + memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer()) require.NoError(t, err) _, err = memory.Append("user", "a") @@ -141,6 +274,86 @@ func TestMemoryRefusesToRebuildAMissingView(t *testing.T) { require.NoError(t, memory.Close()) require.NoError(t, os.Remove(filepath.Join(root, "view.json"))) - _, err = openMemory(t.Context(), root, clock.Now) + _, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer()) require.ErrorIs(t, err, errNoView) } + +// fakeSummarizer builds each node with a line of fakeLine. It fails each task +// while failing is true, and each task waits for gate when gate is not nil. +type fakeSummarizer struct { + mutex sync.Mutex + gate chan struct{} + failing bool + active int + peak int + tasks []Task +} + +func newFakeSummarizer() *fakeSummarizer { + return &fakeSummarizer{} +} + +func (f *fakeSummarizer) Summarize(ctx context.Context, task Task) (string, error) { + f.mutex.Lock() + f.tasks = append(f.tasks, task) + f.active++ + f.peak = max(f.peak, f.active) + fail := f.failing + f.mutex.Unlock() + + defer func() { + f.mutex.Lock() + f.active-- + f.mutex.Unlock() + }() + + if f.gate != nil { + select { + case <-f.gate: + case <-ctx.Done(): + return "", ctx.Err() + } + } + + if fail { + return "", errors.New("fake failure") + } + + return fakeLine(task), nil +} + +// fakeLine returns a line of nodeBytes bytes or less that names the node and +// starts with the input. A real line of this size makes each merge a model +// call. +func fakeLine(task Task) string { + count := 1 << task.Level + line := fmt.Sprintf("summary %d+%d: %s", task.Index*count, count, strings.Join(task.Input, " ")) + line = strings.ReplaceAll(line, "\n", " ") + for len(line) > nodeBytes-12 { + _, size := utf8.DecodeLastRuneInString(line) + line = line[:len(line)-size] + } + + return line +} + +// longText returns a text that is too long for its own node. +func longText(n int) string { + return fmt.Sprintf("long %d ", n) + strings.Repeat("x", nodeBytes) +} + +// SetFailing makes the next tasks fail or not. +func (f *fakeSummarizer) SetFailing(failing bool) { + f.mutex.Lock() + defer f.mutex.Unlock() + + f.failing = failing +} + +// Tasks returns the tasks so far. +func (f *fakeSummarizer) Tasks() []Task { + f.mutex.Lock() + defer f.mutex.Unlock() + + return slices.Clone(f.tasks) +} diff --git a/queue.go b/queue.go new file mode 100644 index 0000000..5f2fcb2 --- /dev/null +++ b/queue.go @@ -0,0 +1,116 @@ +package main + +import ( + "slices" +) + +const ( + // maxBuilds is the number of model builds that can run at once. + maxBuilds = 8 + + // leafWindow is the number of unbuilt leaves that can be before a leaf + // whose build starts. + leafWindow = 8 +) + +// queue holds the nodes that wait for a model build. A leaf joins it when its +// message is too long to be its own node. A merge joins it when both of its +// children are built and their texts do not fit in one node. So the core never +// scans the tree for work. +// +// A leaf can start when fewer than leafWindow leaves before it are unbuilt. +// Leaves start before merges, oldest first, because a turn waits for the leaves +// of all earlier messages. Merges start in the order in which they became +// ready. A failed build waits until the next message. +type queue struct { + leaves []int + merges []nodeKey + queued map[nodeKey]bool + running map[nodeKey]bool + failed map[nodeKey]bool +} + +func newQueue() *queue { + return &queue{ + queued: make(map[nodeKey]bool), + running: make(map[nodeKey]bool), + failed: make(map[nodeKey]bool), + } +} + +// AddLeaf adds the leaf of message id. The IDs must come in increasing order. +func (q *queue) AddLeaf(id int) { + q.leaves = append(q.leaves, id) +} + +// AddMerge adds a merge whose children are built. A merge that is in the queue +// already stays as it is. +func (q *queue) AddMerge(key nodeKey) { + if q.queued[key] { + return + } + + q.queued[key] = true + q.merges = append(q.merges, key) +} + +// Next returns the next node that can start, and marks it as running. +func (q *queue) Next() (nodeKey, bool) { + if len(q.running) >= maxBuilds { + return nodeKey{}, false + } + + for _, id := range q.leaves[:min(len(q.leaves), leafWindow)] { + key := nodeKey{0, id} + if q.startable(key) { + q.running[key] = true + return key, true + } + } + + for _, key := range q.merges { + if q.startable(key) { + q.running[key] = true + return key, true + } + } + + return nodeKey{}, false +} + +func (q *queue) startable(key nodeKey) bool { + return !q.running[key] && !q.failed[key] +} + +// Done removes a built node from the queue. +func (q *queue) Done(key nodeKey) { + delete(q.running, key) + if key.level == 0 { + q.leaves = slices.DeleteFunc(q.leaves, func(id int) bool { + return id == key.index + }) + + return + } + + delete(q.queued, key) + q.merges = slices.DeleteFunc(q.merges, func(other nodeKey) bool { + return other == key + }) +} + +// Fail keeps a node in the queue, but it cannot start until Retry. +func (q *queue) Fail(key nodeKey) { + delete(q.running, key) + q.failed[key] = true +} + +// Retry lets the failed nodes start again. +func (q *queue) Retry() { + clear(q.failed) +} + +// Running returns the number of builds that run. +func (q *queue) Running() int { + return len(q.running) +} diff --git a/tree.go b/tree.go index ee3ff68..30a40ae 100644 --- a/tree.go +++ b/tree.go @@ -1,12 +1,14 @@ package main import ( + "cmp" "context" "encoding/json" "fmt" "math/bits" "os" "path/filepath" + "slices" "time" "unicode/utf8" ) @@ -147,6 +149,20 @@ func (t *Tree) Node(level, index int) (Node, bool) { return node, built } +// Nodes returns the built nodes, ordered by level and then by index. +func (t *Tree) Nodes() []Node { + nodes := make([]Node, 0, len(t.nodes)) + for _, node := range t.nodes { + nodes = append(nodes, node) + } + + slices.SortFunc(nodes, func(a, b Node) int { + return cmp.Or(cmp.Compare(a.Level, b.Level), cmp.Compare(a.Index, b.Index)) + }) + + return nodes +} + // leafSource returns the text that the leaf of message holds or summarizes. func leafSource(message Message) string { return message.Kind + ": " + message.Text