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 eb28341f69821e6158ad49e29fd2b644314ecff3..9c8073f9b06fad474965ec2b9fb78812e0c9e2c9 100644 (file)
--- a/memory.go
+++ b/memory.go
@@ -7,6 +7,7 @@ import (
        "fmt"
        "os"
        "path/filepath"
        "fmt"
        "os"
        "path/filepath"

+       "strings"

        "sync"
        "time"
 )
        "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")
 
 // 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
 // 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 {
 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.
 }
 
 // viewFile is the content of view.json.


@@ -37,19 +63,25 @@ type viewFile struct {
        Compaction viewState `json:"compaction"`
 }
 
        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{
        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())
        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())
        }
 
                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
 }
 
        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 {
 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 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() {
        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
                }
                if err != nil {
                        return err
                }


@@ -121,8 +205,12 @@ func (m *Memory) load() error {
        return m.save()
 }
 
        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 {
 func (m *Memory) Close() error {

+       m.cancel()
+       m.workers.Wait()
+

        m.mutex.Lock()
        defer m.mutex.Unlock()
 
        m.mutex.Lock()
        defer m.mutex.Unlock()
 


@@ -134,7 +222,8 @@ func (m *Memory) Close() error {
        return errors.Join(err, m.log.Close())
 }
 
        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()
 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
        }
 
                return Message{}, err
        }
 

+       m.queue.Retry()
+       m.dispatch()
+

        return message, nil
 }
 
        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 {
 func (m *Memory) addMessage(message Message) error {

-       _, built := m.tree.Node(0, message.ID)

        leaf, free := freeLeaf(message)
        leaf, free := freeLeaf(message)

-       if free && !built {
+       if free {

                err := m.build(leaf)
                if err != nil {
                        return err
                }
                err := m.build(leaf)
                if err != nil {
                        return err
                }

+       } else {
+               m.queue.AddLeaf(message.ID)

        }
 
        merged := m.chat.Append(m.tree)
        }
 
        merged := m.chat.Append(m.tree)


@@ -178,17 +271,23 @@ func (m *Memory) addMessage(message Message) error {
        return nil
 }
 
        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 {
 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)
                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
                }
 
                        return nil
                }
 


@@ -199,13 +298,104 @@ func (m *Memory) build(node Node) error {
 
                parent, free := freeMerge(left, right)
                if !free {
 
                parent, free := freeMerge(left, right)
                if !free {

+                       m.queue.AddMerge(nodeKey{node.Level + 1, node.Index / 2})

                        return nil
                }
 
                        return nil
                }
 

+               err := m.tree.Add(parent)
+               if err != nil {
+                       return err
+               }
+

                node = parent
        }
 }
 
                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{
 // 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 426654a63667c27508eabc2b7f3ad3c31fb4da78..e8682205933daa7feefea465948b999f3f6862f4 100644 (file)
--- a/memory_test.go
+++ b/memory_test.go
@@ -1,10 +1,17 @@
 package main
 
 import (
 package main
 
 import (

+       "context"
+       "errors"

        "fmt"
        "fmt"

+       "maps"

        "os"
        "path/filepath"
        "os"
        "path/filepath"

+       "slices"
+       "strings"
+       "sync"

        "testing"
        "testing"

+       "unicode/utf8"

 
        "github.com/stretchr/testify/require"
 )
 
        "github.com/stretchr/testify/require"
 )


@@ -12,7 +19,7 @@ import (
 func TestMemoryRestartLoadsTheSameView(t *testing.T) {
        root := t.TempDir()
        clock := newTestClock()
 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
        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())
 
        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() {
        require.NoError(t, err)
 
        defer func() {


@@ -47,7 +54,7 @@ func TestMemoryRestartLoadsTheSameView(t *testing.T) {
 }
 
 func TestMemoryKeepsTheCompactionViewCoarser(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() {
        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, "<chat>\n0+1|"+left.Text+"\n1+1|"+right.Text+"\n</chat>", 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) {
 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() {
        require.NoError(t, err)
 
        defer func() {


@@ -94,23 +218,32 @@ func TestMemoryBuildsFreeParents(t *testing.T) {
 }
 
 func TestMemoryKeepsALongMessageOutOfTheView(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())
        }()
 
        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, err)

-
+       require.NoError(t, memory.Wait(t.Context()))

        require.Equal(t, "<chat>\n0+1|"+pending+"\n</chat>", memory.View())
        require.Equal(t, "<chat>\n0+1|"+pending+"\n</chat>", memory.View())

+
+       summarizer.SetFailing(false)
+
+       _, err = memory.Append("user", "again")
+       require.NoError(t, err)
+       require.NoError(t, memory.Wait(t.Context()))
+       require.Equal(t, "<chat>\n0+1|"+fakeLine(memory.task(nodeKey{0, 0}))+"\n1+1|user: again\n</chat>", memory.View())

 }
 
 func TestMemoryCoversMessagesAfterTheSavedView(t *testing.T) {
        root := t.TempDir()
        clock := newTestClock()
 }
 
 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")
        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())
 
        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() {
        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()
 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")
        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")))
 
        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)
 }
        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 (file)
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 ee3ff6807dbe8c660b8aba4cec90e6eb095c9867..30a40ae04aa1737ac782e10a62bf8ec5bf3b87dd 100644 (file)
--- a/tree.go
+++ b/tree.go
@@ -1,12 +1,14 @@
 package main
 
 import (
 package main
 
 import (

+       "cmp"

        "context"
        "encoding/json"
        "fmt"
        "math/bits"
        "os"
        "path/filepath"
        "context"
        "encoding/json"
        "fmt"
        "math/bits"
        "os"
        "path/filepath"

+       "slices"

        "time"
        "unicode/utf8"
 )
        "time"
        "unicode/utf8"
 )


@@ -147,6 +149,20 @@ func (t *Tree) Node(level, index int) (Node, bool) {
        return node, built
 }
 
        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
 // leafSource returns the text that the leaf of message holds or summarizes.
 func leafSource(message Message) string {
        return message.Kind + ": " + message.Text