X-Git-Url: https://git.giacomolli.uk/onechat/blobdiff_plain/e5343bd4c3ec69192e92f4be339b855eef36b6ec..522d226230a2eae0a8a33ae0fbf32f96c6886b26:/memory_test.go
diff --git a/memory_test.go b/memory_test.go
index 426654a..8c34149 100644
--- a/memory_test.go
+++ b/memory_test.go
@@ -1,10 +1,18 @@
package main
import (
+ "context"
+ "errors"
"fmt"
+ "maps"
"os"
"path/filepath"
+ "slices"
+ "strings"
+ "sync"
"testing"
+ "time"
+ "unicode/utf8"
"github.com/stretchr/testify/require"
)
@@ -12,7 +20,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 +39,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 +55,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 +83,129 @@ 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)
+
+ 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, isBuilt(t, memory.tree, 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, isBuilt(t, memory.tree, 0, 8))
+ summarizer.SetFailing(false)
+
+ _, err = memory.Append("user", "retry")
+ require.NoError(t, err)
+ require.NoError(t, memory.Wait(t.Context()))
+ require.True(t, isBuilt(t, memory.tree, 0, 8))
+ require.True(t, isBuilt(t, memory.tree, 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 := readNode(t, memory.tree, 0, 0)
+ right := readNode(t, memory.tree, 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, isBuilt(t, memory.tree, 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, isBuilt(t, memory.tree, 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() {
@@ -88,29 +217,40 @@ func TestMemoryBuildsFreeParents(t *testing.T) {
require.NoError(t, err)
}
- node, built := memory.tree.Node(2, 0)
- require.True(t, built)
+ node := readNode(t, memory.tree, 2, 0)
require.Equal(t, "user: a\nuser: b\nuser: c\nuser: d", node.Text)
}
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()))
+
+ task, err := memory.task(nodeKey{0, 0})
+ require.NoError(t, err)
+ require.Equal(t, "\n0+1|"+fakeLine(task)+"\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 +260,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 +273,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 +281,94 @@ 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
+}
+
+// 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)
+}