Remove two timing gaps from the memory tests

[onechat] / memory_test.go

diff --git a/memory_test.go b/memory_test.go
index b5bddb04ebfb78d255133c4b712b7698e2fce8ec..ddbb014030f7cc26bf7f63fc5b128cb047ad05a9 100644 (file)
--- a/memory_test.go
+++ b/memory_test.go
@@ -1,10 +1,18 @@
 package main
 
 import (
 package main
 
 import (

+       "context"
+       "errors"

        "fmt"
        "fmt"

+       "maps"

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

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

        "testing"
        "testing"

+       "time"
+       "unicode/utf8"

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


@@ -12,23 +20,26 @@ 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
        memory.chat.low = 100
        require.NoError(t, err)
 
        memory.chat.high = 200
        memory.chat.low = 100

+       memory.cview.high = 80
+       memory.cview.low = 40

        for n := range 100 {
                _, err = memory.Append("user", fmt.Sprint(n))
                require.NoError(t, err)
        }
 
        view := memory.View()
        for n := range 100 {
                _, err = memory.Append("user", fmt.Sprint(n))
                require.NoError(t, err)
        }
 
        view := memory.View()

+       context := memory.Context(0, 100)

        saved, err := os.ReadFile(filepath.Join(root, "view.json"))
        require.NoError(t, err)
        require.Less(t, len(memory.chat.lines), 100)
        require.NoError(t, memory.Close())
 
        saved, err := os.ReadFile(filepath.Join(root, "view.json"))
        require.NoError(t, err)
        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() {


@@ -36,14 +47,165 @@ func TestMemoryRestartLoadsTheSameView(t *testing.T) {
        }()
 
        require.Equal(t, view, memory.View())
        }()
 
        require.Equal(t, view, memory.View())

+       require.Equal(t, context, memory.Context(0, 100))

 
        loaded, err := os.ReadFile(filepath.Join(root, "view.json"))
        require.NoError(t, err)
        require.Equal(t, saved, loaded)
 }
 
 
        loaded, err := os.ReadFile(filepath.Join(root, "view.json"))
        require.NoError(t, err)
        require.Equal(t, saved, loaded)
 }
 

+func TestMemoryKeepsTheCompactionViewCoarser(t *testing.T) {
+       memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, newFakeSummarizer())
+       require.NoError(t, err)
+
+       defer func() {
+               require.NoError(t, memory.Close())
+       }()
+
+       memory.chat.high = 200
+       memory.chat.low = 100
+       memory.cview.high = 80
+       memory.cview.low = 40
+       for n := range 64 {
+               _, err = memory.Append("user", fmt.Sprint(n%10))
+               require.NoError(t, err)
+
+               coarse := 0
+               for _, line := range memory.chat.lines {
+                       outer := memory.cview.lines[coarse]
+                       for nodeEnd(outer.level, outer.index) <= line.index<<line.level {
+                               coarse++
+                               outer = memory.cview.lines[coarse]
+                       }
+
+                       require.GreaterOrEqual(t, outer.level, line.level, "message %d", n)
+               }
+       }
+}
+
+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)
+
+       require.Eventually(t, func() bool {
+               return summarizer.Active() == maxBuilds
+       }, 10*time.Second, time.Millisecond)
+
+       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() {


@@ -61,23 +223,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")


@@ -87,7 +258,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() {


@@ -100,7 +271,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")


@@ -108,6 +279,94 @@ 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
+}
+
+// Active returns the number of tasks that run now.
+func (f *fakeSummarizer) Active() int {
+       f.mutex.Lock()
+       defer f.mutex.Unlock()
+
+       return f.active
+}
+
+// Tasks returns the tasks so far.
+func (f *fakeSummarizer) Tasks() []Task {
+       f.mutex.Lock()
+       defer f.mutex.Unlock()
+
+       return slices.Clone(f.tasks)
+}