func newBuildQueue() *BuildQueue {
return &BuildQueue{
- queued: make(map[nodeKey]bool),
- running: make(map[nodeKey]bool),
- failed: make(map[nodeKey]bool),
+ queued: make(map[NodeKey]bool),
+ running: make(map[NodeKey]bool),
+ failed: make(map[NodeKey]bool),
}
}
type BuildQueue struct {
leaves []int
merges BuildMergeHeap
- queued map[nodeKey]bool
- running map[nodeKey]bool
- failed map[nodeKey]bool
- retry []nodeKey
+ queued map[NodeKey]bool
+ running map[NodeKey]bool
+ failed map[NodeKey]bool
+ retry []NodeKey
}
// AddLeaf adds the leaf of message id. The IDs must come in increasing order.
// AddMerge adds a merge whose children are built. A merge that is in the queue
// already stays as it is.
-func (q *BuildQueue) AddMerge(key nodeKey) {
+func (q *BuildQueue) AddMerge(key NodeKey) {
if q.queued[key] {
return
}
}
// Next returns the next node that can start, and marks it as running.
-func (q *BuildQueue) Next() (nodeKey, bool) {
+func (q *BuildQueue) Next() (NodeKey, bool) {
if len(q.running) >= buildMax {
- return nodeKey{}, false
+ return NodeKey{}, false
}
leaf, found := q.nextLeaf()
if len(q.merges) > 0 && (!found || nodeEnd(q.merges[0].level, q.merges[0].index) <= leaf.index) {
- key := heap.Pop(&q.merges).(nodeKey)
+ key := heap.Pop(&q.merges).(NodeKey)
q.running[key] = true
return key, true
}
// Done removes a built node from the queue.
-func (q *BuildQueue) Done(key nodeKey) {
+func (q *BuildQueue) Done(key NodeKey) {
delete(q.running, key)
if key.level > 0 {
delete(q.queued, key)
}
// Fail keeps a node in the queue, but it cannot start until Retry.
-func (q *BuildQueue) Fail(key nodeKey) {
+func (q *BuildQueue) Fail(key NodeKey) {
delete(q.running, key)
if key.level > 0 {
q.retry = append(q.retry, key)
}
// nextLeaf returns the oldest leaf in the window that can start.
-func (q *BuildQueue) nextLeaf() (nodeKey, bool) {
+func (q *BuildQueue) nextLeaf() (NodeKey, bool) {
for _, id := range q.leaves[:min(len(q.leaves), buildLeafWindow)] {
- key := nodeKey{0, id}
+ key := NodeKey{0, id}
if !q.running[key] && !q.failed[key] {
return key, true
}
}
- return nodeKey{}, false
+ return NodeKey{}, false
}
// BuildMergeHeap holds the merges that can start, by the end of their context.
-type BuildMergeHeap []nodeKey
+type BuildMergeHeap []NodeKey
func (h BuildMergeHeap) Len() int {
return len(h)
}
func (h *BuildMergeHeap) Push(value any) {
- *h = append(*h, value.(nodeKey))
+ *h = append(*h, value.(NodeKey))
}
func (h *BuildMergeHeap) Pop() any {
queue := newBuildQueue()
queue.AddLeaf(4)
queue.AddLeaf(5)
- queue.AddMerge(nodeKey{1, 2})
- queue.AddMerge(nodeKey{1, 1})
+ queue.AddMerge(NodeKey{1, 2})
+ queue.AddMerge(NodeKey{1, 1})
- var order []nodeKey
+ var order []NodeKey
for {
key, ready := queue.Next()
if !ready {
order = append(order, key)
}
- require.Equal(t, []nodeKey{{1, 1}, {0, 4}, {0, 5}, {1, 2}}, order)
+ require.Equal(t, []NodeKey{{1, 1}, {0, 4}, {0, 5}, {1, 2}}, order)
}
func TestQueueRetriesAFailedMergeAtTheNextMessage(t *testing.T) {
queue := newBuildQueue()
- queue.AddMerge(nodeKey{1, 0})
+ queue.AddMerge(NodeKey{1, 0})
key, ready := queue.Next()
require.True(t, ready)
queue.Retry()
key, ready = queue.Next()
require.True(t, ready)
- require.Equal(t, nodeKey{1, 0}, key)
+ require.Equal(t, NodeKey{1, 0}, key)
queue.Done(key)
- queue.AddMerge(nodeKey{1, 0})
+ queue.AddMerge(NodeKey{1, 0})
_, ready = queue.Next()
require.True(t, ready)
}
import "time"
-// testClock is a clock that tests set by hand.
-type testClock struct {
+// TestClock is a clock that tests set by hand.
+type TestClock struct {
now time.Time
}
-func newTestClock() *testClock {
- return &testClock{
- now: time.Date(2026, 1, 1, 12, 0, 0, 0, time.UTC),
- }
+func (c *TestClock) Now() time.Time {
+ return c.now
}
-func (c *testClock) Now() time.Time {
- return c.now
+func newTestClock() *TestClock {
+ return &TestClock{
+ now: time.Date(2026, 1, 1, 12, 0, 0, 0, time.UTC),
+ }
}
Instructions: instructions,
Tools: c.tools,
Input: []any{modelUserMessage(parts...)},
- Node: &nodeKey{task.Level, task.Index},
+ Node: &NodeKey{task.Level, task.Index},
})
if err != nil {
break
_, err = memory.Append("echo", strings.Repeat("output ", 200))
require.NoError(t, err)
require.NoError(t, memory.Wait(t.Context()))
- require.False(t, isBuilt(t, memory.tree, 0, 0))
+ require.False(t, treeIsBuilt(t, memory.tree, 0, 0))
_, err = memory.Append("user", "next")
require.NoError(t, err)
require.NoError(t, memory.Wait(t.Context()))
- node := readNode(t, memory.tree, 0, 0)
+ node := treeReadNode(t, memory.tree, 0, 0)
require.Equal(t, "summary of the long output", node.Text)
require.Equal(t, "<chat>\n0+1|summary of the long output\n1+1|user: next\n</chat>", memory.View())
require.Equal(t, 2, server.Calls())
"strings"
)
-// position names one line of a day file: the day as YYYYMMDD, the byte offset
+// Position names one line of a day file: the day as YYYYMMDD, the byte offset
// of the line, and its length with the newline.
-type position struct {
+type Position struct {
day uint32
offset uint64
length uint32
}
-// journal appends JSON lines to one file for each day in dir, named
+// Journal appends JSON lines to one file for each day in dir, named
// YYYY-MM-DD.jsonl, and reads them back by their position. It keeps the file of
// the last day open for writes and each file that it read open for reads.
-type journal struct {
+type Journal struct {
dir string
file *os.File
day uint32
// Write appends line to the file of day, YYYY-MM-DD, with one write to the
// operating system, without a buffer. A crash of the process then loses
// nothing. It returns the position of the line.
-func (j *journal) Write(day string, line []byte) (position, error) {
- number, err := dayNumber(day)
+func (j *Journal) Write(day string, line []byte) (Position, error) {
+ number, err := journalDayNumber(day)
if err != nil {
- return position{}, err
+ return Position{}, err
}
if number != j.day {
err = j.closeWriter()
if err != nil {
- return position{}, err
+ return Position{}, err
}
j.file, err = os.OpenFile(filepath.Join(j.dir, day+".jsonl"), os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o600)
if err != nil {
- return position{}, fmt.Errorf("open day file: %w", err)
+ return Position{}, fmt.Errorf("open day file: %w", err)
}
info, err := j.file.Stat()
if err != nil {
- return position{}, fmt.Errorf("read day file size: %w", err)
+ return Position{}, fmt.Errorf("read day file size: %w", err)
}
j.day = number
_, err = j.file.Write(line)
if err != nil {
- return position{}, fmt.Errorf("write line: %w", err)
+ return Position{}, fmt.Errorf("write line: %w", err)
}
- place := position{
+ place := Position{
day: number,
offset: j.size,
length: uint32(len(line)),
return place, nil
}
-func (j *journal) closeWriter() error {
- if j.file == nil {
- return nil
- }
-
- err := j.file.Close()
- j.file = nil
- j.day = 0
- if err != nil {
- return fmt.Errorf("close day file: %w", err)
- }
-
- return nil
-}
-
// Read returns the line at place without its newline.
-func (j *journal) Read(place position) ([]byte, error) {
+func (j *Journal) Read(place Position) ([]byte, error) {
reader, open := j.readers[place.day]
if !open {
var err error
- reader, err = os.Open(filepath.Join(j.dir, dayName(place.day)+".jsonl"))
+ reader, err = os.Open(filepath.Join(j.dir, journalDayName(place.day)+".jsonl"))
if err != nil {
return nil, fmt.Errorf("open day file: %w", err)
}
line := make([]byte, place.length)
_, err := reader.ReadAt(line, int64(place.offset))
if err != nil {
- return nil, fmt.Errorf("read line of %s: %w", dayName(place.day), err)
+ return nil, fmt.Errorf("read line of %s: %w", journalDayName(place.day), err)
}
if len(line) == 0 || line[len(line)-1] != '\n' {
- return nil, fmt.Errorf("line of %s at %d does not end its record", dayName(place.day), place.offset)
+ return nil, fmt.Errorf("line of %s at %d does not end its record", journalDayName(place.day), place.offset)
}
return line[:len(line)-1], nil
}
// Close closes the open day files.
-func (j *journal) Close() error {
+func (j *Journal) Close() error {
err := j.closeWriter()
for day, reader := range j.readers {
err = errors.Join(err, reader.Close())
// position, oldest day first. The zero position starts at the first line. A
// crash during a write can leave a part of a line at the end of the last day
// file. Its write never returned, so Scan cuts it off.
-func (j *journal) Scan(ctx context.Context, place position, add func([]byte, position) error) error {
+func (j *Journal) Scan(ctx context.Context, place Position, add func([]byte, Position) error) error {
paths, err := filepath.Glob(filepath.Join(j.dir, "*.jsonl"))
if err != nil {
return fmt.Errorf("list day files: %w", err)
slices.Sort(paths)
for n, path := range paths {
- day, err := dayNumber(strings.TrimSuffix(filepath.Base(path), ".jsonl"))
+ day, err := journalDayNumber(strings.TrimSuffix(filepath.Base(path), ".jsonl"))
if err != nil {
return err
}
offset = place.offset
}
- err = scanFile(ctx, path, day, offset, n == len(paths)-1, add)
+ err = journalScanFile(ctx, path, day, offset, n == len(paths)-1, add)
if err != nil {
return err
}
return nil
}
-// scanFile gives each line of the day file at path from offset on to add. In
-// the last day file, it cuts a part of a line at the end.
-func scanFile(ctx context.Context, path string, day uint32, offset uint64, last bool, add func([]byte, position) error) (err error) {
+func (j *Journal) closeWriter() error {
+ if j.file == nil {
+ return nil
+ }
+
+ err := j.file.Close()
+ j.file = nil
+ j.day = 0
+ if err != nil {
+ return fmt.Errorf("close day file: %w", err)
+ }
+
+ return nil
+}
+
+// journalDayNumber returns the day YYYY-MM-DD as the number YYYYMMDD.
+func journalDayNumber(day string) (uint32, error) {
+ number, err := strconv.ParseUint(strings.ReplaceAll(day, "-", ""), 10, 32)
+ if err != nil || len(day) != len("2006-01-02") {
+ return 0, fmt.Errorf("invalid day %q", day)
+ }
+
+ return uint32(number), nil
+}
+
+// journalDayName returns the day number YYYYMMDD as YYYY-MM-DD.
+func journalDayName(day uint32) string {
+ return fmt.Sprintf("%04d-%02d-%02d", day/10000, day/100%100, day%100)
+}
+
+// journalScanFile gives each line of the day file at path from offset on to
+// add. In the last day file, it cuts a part of a line at the end.
+func journalScanFile(ctx context.Context, path string, day uint32, offset uint64, last bool, add func([]byte, Position) error) (err error) {
file, err := os.Open(path)
if err != nil {
return fmt.Errorf("open day file: %w", err)
}
if errors.Is(err, io.EOF) && last {
- return cutFile(path, offset)
+ return journalCutFile(path, offset)
}
if errors.Is(err, io.EOF) {
return fmt.Errorf("read %s: %w", filepath.Base(path), err)
}
- place := position{
+ place := Position{
day: day,
offset: offset,
length: uint32(len(line)),
}
}
-// cutFile cuts the file at path to size bytes.
-func cutFile(path string, size uint64) error {
+// journalCutFile cuts the file at path to size bytes.
+func journalCutFile(path string, size uint64) error {
err := os.Truncate(path, int64(size))
if err != nil {
return fmt.Errorf("cut incomplete line: %w", err)
return nil
}
-
-// dayNumber returns the day YYYY-MM-DD as the number YYYYMMDD.
-func dayNumber(day string) (uint32, error) {
- number, err := strconv.ParseUint(strings.ReplaceAll(day, "-", ""), 10, 32)
- if err != nil || len(day) != len("2006-01-02") {
- return 0, fmt.Errorf("invalid day %q", day)
- }
-
- return uint32(number), nil
-}
-
-// dayName returns the day number YYYYMMDD as YYYY-MM-DD.
-func dayName(day uint32) string {
- return fmt.Sprintf("%04d-%02d-%02d", day/10000, day/100%100, day%100)
-}
"unicode/utf8"
)
-// kinds holds the kinds that a message can have.
-var kinds = map[string]bool{
+// logKinds holds the logKinds that a message can have.
+var logKinds = map[string]bool{
"user": true,
"reply": true,
"tool": true,
"note": true,
}
-// Message is one entry of the chat log. ID is the permanent position of the
-// message in the chat, from 0. Size is the UTF-8 byte length of Text. Date is
-// the time of the write in RFC 3339 format.
-type Message struct {
- ID int `json:"i"`
- Kind string `json:"kind"`
- Text string `json:"text"`
- Size int `json:"size"`
- Date string `json:"date"`
-}
-
-// Log is the append-only message log of one chat. It writes one JSON line for
-// each message to main/YYYY-MM-DD.jsonl, named for the local day of the write.
-// It never edits or deletes a line. Each line goes to the operating system
-// before Append returns, so a crash of the process loses no message. An
-// exclusive lock on the lock file makes sure that only one process writes the
-// log. After a failed write, the end of the file is not known, so the log
-// refuses all writes until it is opened again. The caller serializes all calls.
-//
-// The log keeps no message in memory. The index index/main.idx holds the
-// position of message i in record i, so a read of a message is one read of the
-// index and one of its day file. At open the log indexes the lines after the
-// last indexed one, so a missing index is built again from the day files.
-type Log struct {
- clock func() time.Time
- lock *os.File
- journal journal
- index *records
- failure error
-}
-
// openLog locks the chat in root and indexes its new lines. The clock gives the
// date of each new message.
func openLog(ctx context.Context, root string, clock func() time.Time) (*Log, error) {
l := &Log{
clock: clock,
lock: lock,
- journal: journal{
+ journal: Journal{
dir: filepath.Join(root, "main"),
},
}
return l, nil
}
-// catchUp indexes the lines after the last indexed line.
-func (l *Log) catchUp(ctx context.Context) error {
- start := position{}
- if l.index.count > 0 {
- record, _, err := l.index.Read(l.index.count - 1)
- if err != nil {
- return err
- }
-
- last := decodePosition(record)
- start = position{
- day: last.day,
- offset: last.offset + uint64(last.length),
- }
- }
-
- return l.journal.Scan(ctx, start, func(line []byte, place position) error {
- var message Message
- err := json.Unmarshal(line, &message)
- if err != nil {
- return err
- }
-
- if message.ID != l.index.count {
- return fmt.Errorf("expected message %d, found %d", l.index.count, message.ID)
- }
-
- err = validateMessage(message)
- if err != nil {
- return fmt.Errorf("message %d: %w", message.ID, err)
- }
-
- return l.addIndex(place)
- })
-}
-
-func (l *Log) addIndex(place position) error {
- record := make([]byte, positionBytes)
- encodePosition(record, place)
-
- return l.index.Append(record)
+// Log is the append-only message log of one chat. It writes one JSON line for
+// each message to main/YYYY-MM-DD.jsonl, named for the local day of the write.
+// It never edits or deletes a line. Each line goes to the operating system
+// before Append returns, so a crash of the process loses no message. An
+// exclusive lock on the lock file makes sure that only one process writes the
+// log. After a failed write, the end of the file is not known, so the log
+// refuses all writes until it is opened again. The caller serializes all calls.
+//
+// The log keeps no message in memory. The index index/main.idx holds the
+// position of message i in record i, so a read of a message is one read of the
+// index and one of its day file. At open the log indexes the lines after the
+// last indexed one, so a missing index is built again from the day files.
+type Log struct {
+ clock func() time.Time
+ lock *os.File
+ journal Journal
+ index *Records
+ failure error
}
// Close releases the lock. The log refuses all calls after Close.
Date: now.Format(time.RFC3339Nano),
}
- err := validateMessage(message)
+ err := logValidateMessage(message)
if err != nil {
return Message{}, err
}
return message, nil
}
-// validateMessage makes sure that a message has a known kind, valid UTF-8 text,
-// the size of its text, and a valid date.
-func validateMessage(message Message) error {
- if !kinds[message.Kind] {
+// catchUp indexes the lines after the last indexed line.
+func (l *Log) catchUp(ctx context.Context) error {
+ start := Position{}
+ if l.index.count > 0 {
+ record, _, err := l.index.Read(l.index.count - 1)
+ if err != nil {
+ return err
+ }
+
+ last := decodePosition(record)
+ start = Position{
+ day: last.day,
+ offset: last.offset + uint64(last.length),
+ }
+ }
+
+ return l.journal.Scan(ctx, start, func(line []byte, place Position) error {
+ var message Message
+ err := json.Unmarshal(line, &message)
+ if err != nil {
+ return err
+ }
+
+ if message.ID != l.index.count {
+ return fmt.Errorf("expected message %d, found %d", l.index.count, message.ID)
+ }
+
+ err = logValidateMessage(message)
+ if err != nil {
+ return fmt.Errorf("message %d: %w", message.ID, err)
+ }
+
+ return l.addIndex(place)
+ })
+}
+
+func (l *Log) addIndex(place Position) error {
+ record := make([]byte, positionBytes)
+ encodePosition(record, place)
+
+ return l.index.Append(record)
+}
+
+// Message is one entry of the chat log. ID is the permanent position of the
+// message in the chat, from 0. Size is the UTF-8 byte length of Text. Date is
+// the time of the write in RFC 3339 format.
+type Message struct {
+ ID int `json:"i"`
+ Kind string `json:"kind"`
+ Text string `json:"text"`
+ Size int `json:"size"`
+ Date string `json:"date"`
+}
+
+// logValidateMessage makes sure that a message has a known kind, valid UTF-8
+// text, the size of its text, and a valid date.
+func logValidateMessage(message Message) error {
+ if !logKinds[message.Kind] {
return fmt.Errorf("unknown message kind %q", message.Kind)
}
require.NoError(t, log.Close())
// A missing index, and an index that a crash cut in a record.
- for _, cut := range []func(string) error{os.Remove, truncateToOneAndAHalf} {
+ for _, cut := range []func(string) error{os.Remove, logTruncateToOneAndAHalf} {
require.NoError(t, cut(filepath.Join(root, "index", "main.idx")))
log, err = openLog(t.Context(), root, clock.Now)
require.ErrorContains(t, err, "ends in an incomplete line")
}
-// truncateToOneAndAHalf cuts the index at path in its second record.
-func truncateToOneAndAHalf(path string) error {
+// logTruncateToOneAndAHalf cuts the index at path in its second record.
+func logTruncateToOneAndAHalf(path string) error {
return os.Truncate(path, positionBytes+positionBytes/2)
}
"time"
)
-// errNoView tells that a chat has messages but no saved view.
-var errNoView = errors.New("the chat has messages but no saved view")
+// errMemoryNoView tells that a chat has messages but no saved view.
+var errMemoryNoView = errors.New("the chat has messages but no saved view")
+
+// 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,
+ queue: newBuildQueue(),
+ summarizer: summarizer,
+ 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)
+ if err != nil {
+ return nil, errors.Join(err, m.Close())
+ }
+
+ err = m.load()
+ if err != nil {
+ 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
+}
// 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, 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
mutex sync.Mutex
root string
log *Log
- syncer *syncer
+ syncer *Syncer
tree *Tree
chat *View
cview *View
changed chan struct{}
}
-// viewFile is the content of view.json.
-type viewFile struct {
- Chat ViewState `json:"chat"`
- Compaction ViewState `json:"compaction"`
-}
+// 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()
-// 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.mutex.Lock()
+ defer m.mutex.Unlock()
+
+ var err error
+ if m.syncer != nil {
+ err = m.syncer.Close()
}
- m := &Memory{
- root: root,
- log: log,
- queue: newBuildQueue(),
- summarizer: summarizer,
- changed: make(chan struct{}),
+ if m.tree != nil {
+ err = errors.Join(err, m.tree.Close())
}
- m.ctx, m.cancel = context.WithCancel(ctx)
+ return errors.Join(err, m.log.Close())
+}
- m.syncer, err = openSyncer(root)
+// 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()
+
+ message, err := m.log.Append(kind, text)
if err != nil {
- return nil, errors.Join(err, m.Close())
+ return Message{}, err
}
- m.tree, err = openTree(ctx, root, log.Len(), clock)
+ defer m.syncer.Written()
+
+ err = m.addMessage(message)
if err != nil {
- return nil, errors.Join(err, m.Close())
+ return Message{}, err
}
- err = m.load()
+ err = m.save()
if err != nil {
- return nil, errors.Join(fmt.Errorf("load view: %w", err), m.Close())
+ return Message{}, err
}
- err = m.restoreQueue()
- if err != nil {
- return nil, errors.Join(fmt.Errorf("restore build queue: %w", err), m.Close())
+ m.queue.Retry()
+ m.dispatch()
+
+ 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
}
- err = m.catchUp()
- if err != nil {
- return nil, errors.Join(err, m.Close())
+ 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()
+ }
}
+}
+// Flush returns when every write before the call is durable.
+func (m *Memory) Flush(ctx context.Context) error {
+ return m.syncer.Flush(ctx)
+}
+
+// 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()
- m.syncer.Written()
- m.dispatch()
- m.mutex.Unlock()
+ defer m.mutex.Unlock()
- return m, nil
+ 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.
}
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)
return m.save()
}
-// 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) {
- m.mutex.Lock()
- defer m.mutex.Unlock()
-
- message, err := m.log.Append(kind, text)
- if err != nil {
- return Message{}, err
- }
-
- defer m.syncer.Written()
+// 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
+ }
- err = m.addMessage(message)
- if err != nil {
- return Message{}, err
- }
+ task, err := m.task(key)
+ if err != nil {
+ m.queue.Fail(key)
+ continue
+ }
- err = m.save()
- if err != nil {
- return Message{}, err
+ m.workers.Add(1)
+ go m.run(key, task)
}
-
- m.queue.Retry()
- m.dispatch()
-
- return message, nil
}
// addMessage builds the leaf of message when it needs no model call, or else
return nil
}
+// 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)
+ }
+
+ err = writeFileAtomic(m.viewPath(), append(data, '\n'))
+ if err != nil {
+ return fmt.Errorf("save view: %w", err)
+ }
+
+ return nil
+}
+
+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 {
m.chat.Refresh(node)
m.cview.Refresh(node)
- return m.settle(nodeKey{node.Level, node.Index})
+ 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 {
+func (m *Memory) settle(key NodeKey) error {
for {
- parent := nodeKey{key.level + 1, key.index / 2}
+ parent := NodeKey{key.level + 1, key.index / 2}
ready, err := m.ready(parent)
if err != nil || !ready {
return 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
-}
-
-// 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
-}
-
-// 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
- }
-
- task, err := m.task(key)
- if err != nil {
- m.queue.Fail(key)
- continue
- }
-
- m.workers.Add(1)
- go m.run(key, task)
- }
-}
-
// 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) {
+func (m *Memory) task(key NodeKey) (Task, error) {
task := Task{
Level: key.level,
Index: key.index,
// 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) {
+func (m *Memory) run(key NodeKey, task Task) {
defer m.workers.Done()
text, err := m.summarizer.Summarize(m.ctx, task)
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()
+// 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
}
}
-}
-// Flush returns when every write before the call is durable.
-func (m *Memory) Flush(ctx context.Context) error {
- return m.syncer.Flush(ctx)
-}
+ built, err := m.tree.Built(key.level, key.index)
-// 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(),
- })
- if err != nil {
- return fmt.Errorf("encode view: %w", err)
- }
+ return !built, err
+}
- err = writeFileAtomic(m.viewPath(), append(data, '\n'))
- if err != nil {
- return fmt.Errorf("save view: %w", err)
+func (m *Memory) context(level, index int) string {
+ end := index
+ if level > 0 {
+ end = nodeEnd(level, index)
}
- return nil
-}
-
-func (m *Memory) viewPath() string {
- return filepath.Join(m.root, "view.json")
+ return m.cview.Context(end)
}
// View returns the chat view as the model reads it.
return m.chat.Render()
}
-// 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)
-}
-
// Context returns the compaction view for the build of node (l, i). It ends
// before the message of a leaf, or after the last message of a merge, and it
// stops at the first line that is not built.
return m.context(level, index)
}
-func (m *Memory) context(level, index int) string {
- end := index
- if level > 0 {
- end = nodeEnd(level, index)
- }
+// 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
+}
- return m.cview.Context(end)
+// MemoryViewFile is the content of view.json.
+type MemoryViewFile struct {
+ Chat ViewState `json:"chat"`
+ Compaction ViewState `json:"compaction"`
}
package main
import (
- "context"
- "errors"
"fmt"
"maps"
"os"
"path/filepath"
"slices"
"strings"
- "sync"
"testing"
"time"
- "unicode/utf8"
"github.com/stretchr/testify/require"
)
}()
for n := range 12 {
- _, err = memory.Append("echo", longText(n))
+ _, err = memory.Append("echo", memoryLongText(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 {
+ 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.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() == buildMax
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)
+ require.True(t, treeIsBuilt(t, memory.tree, level, index), "node (%d, %d)", level, index)
}
}
}
}()
for n := range 9 {
- _, err = memory.Append("echo", longText(n))
+ _, err = memory.Append("echo", memoryLongText(n))
require.NoError(t, err)
}
require.Less(t, task.Index, 8)
}
- require.False(t, isBuilt(t, memory.tree, 0, 8))
+ require.False(t, treeIsBuilt(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))
+ require.True(t, treeIsBuilt(t, memory.tree, 0, 8))
+ require.True(t, treeIsBuilt(t, memory.tree, 3, 0))
}
func TestMemoryMergesAfterBothHalves(t *testing.T) {
require.NoError(t, memory.Close())
}()
- _, err = memory.Append("echo", longText(0))
+ _, err = memory.Append("echo", memoryLongText(0))
require.NoError(t, err)
require.NoError(t, memory.Wait(t.Context()))
- _, err = memory.Append("echo", longText(1))
+ _, err = memory.Append("echo", memoryLongText(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)
+ left := treeReadNode(t, memory.tree, 0, 0)
+ right := treeReadNode(t, memory.tree, 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, isBuilt(t, memory.tree, 1, 0))
+ require.True(t, treeIsBuilt(t, memory.tree, 1, 0))
}
func TestMemoryRestartRestoresTheQueue(t *testing.T) {
require.NoError(t, err)
for n := range 4 {
- _, err = memory.Append("echo", longText(n))
+ _, err = memory.Append("echo", memoryLongText(n))
require.NoError(t, err)
}
}()
require.NoError(t, memory.Wait(t.Context()))
- require.True(t, isBuilt(t, memory.tree, 2, 0))
+ require.True(t, treeIsBuilt(t, memory.tree, 2, 0))
}
func TestMemoryBuildsFreeParents(t *testing.T) {
require.NoError(t, err)
}
- node := readNode(t, memory.tree, 2, 0)
+ node := treeReadNode(t, memory.tree, 2, 0)
require.Equal(t, "user: a\nuser: b\nuser: c\nuser: d", node.Text)
}
require.NoError(t, memory.Close())
}()
- _, err = memory.Append("echo", longText(0))
+ _, err = memory.Append("echo", memoryLongText(0))
require.NoError(t, err)
require.NoError(t, memory.Wait(t.Context()))
require.Equal(t, "<chat>\n0+1|"+viewPending+"\n</chat>", memory.View())
require.NoError(t, err)
require.NoError(t, memory.Wait(t.Context()))
- task, err := memory.task(nodeKey{0, 0})
+ task, err := memory.task(NodeKey{0, 0})
require.NoError(t, err)
require.Equal(t, "<chat>\n0+1|"+fakeLine(task)+"\n1+1|user: again\n</chat>", memory.View())
}
require.NoError(t, os.Remove(filepath.Join(root, "view.json")))
_, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
- require.ErrorIs(t, err, errNoView)
+ require.ErrorIs(t, err, errMemoryNoView)
}
-// 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 {
+// memoryLongText returns a text that is too long for its own node.
+func memoryLongText(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)
-}
Instructions string
Tools []ModelTool
Input []any
- Node *nodeKey
+ Node *NodeKey
Turn string
}
client := newTestModelClient(t, server)
request := modelTestRequest()
- request.Node = &nodeKey{2, 5}
+ request.Node = &NodeKey{2, 5}
_, err := client.Create(t.Context(), request)
require.NoError(t, err)
require.NotContains(t, server.Call(0).body, "node")
// probeBuiltByModel 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 probeBuiltByModel(memory *Memory, key nodeKey) (bool, error) {
+func probeBuiltByModel(memory *Memory, key NodeKey) (bool, error) {
if key.level == 0 {
message, err := memory.log.Message(key.index)
if err != nil {
long := 0
for _, message := range messages {
counts[message.Kind]++
- require.True(t, kinds[message.Kind], message.Kind)
+ require.True(t, logKinds[message.Kind], message.Kind)
require.NotEmpty(t, message.Text)
for _, char := range message.Text {
require.True(t, char == '\n' || (char >= ' ' && char <= '~'), "character %q", char)
}
// probeLog returns count messages of made-up work on a small bakery order
-// system, with the mix of kinds and sizes of syntheticMessage. The same count
+// system, with the mix of kinds and sizes of replayMessage. The same count
// always gives the same log.
func probeLog(count int) []ProbeEntry {
g := ProbeWriter{
// positionBytes is the size of an encoded position.
const positionBytes = 16
-// records is a file of fixed-size records, numbered from 0. Record n starts at
-// byte n*size. A record that was never written reads as zeros. An index is data
-// that the day files can give again, so records cuts a partial record at the
-// end of the file.
-type records struct {
- file *os.File
- size int
- count int
-}
-
// openRecords opens or creates the record file at path with records of size
// bytes.
-func openRecords(path string, size int) (*records, error) {
+func openRecords(path string, size int) (*Records, error) {
file, err := os.OpenFile(path, os.O_CREATE|os.O_RDWR, 0o600)
if err != nil {
return nil, fmt.Errorf("open index: %w", err)
return nil, errors.Join(fmt.Errorf("read index size: %w", err), file.Close())
}
- r := &records{
+ r := &Records{
file: file,
size: size,
count: int(info.Size()) / size,
return r, nil
}
+// encodePosition writes place into the first positionBytes bytes of record.
+func encodePosition(record []byte, place Position) {
+ binary.LittleEndian.PutUint32(record[0:], place.day)
+ binary.LittleEndian.PutUint32(record[4:], place.length)
+ binary.LittleEndian.PutUint64(record[8:], place.offset)
+}
+
+// decodePosition reads a position from the first positionBytes bytes of record.
+func decodePosition(record []byte) Position {
+ return Position{
+ day: binary.LittleEndian.Uint32(record[0:]),
+ length: binary.LittleEndian.Uint32(record[4:]),
+ offset: binary.LittleEndian.Uint64(record[8:]),
+ }
+}
+
+// Records is a file of fixed-size Records, numbered from 0. Record n starts at
+// byte n*size. A record that was never written reads as zeros. An index is data
+// that the day files can give again, so Records cuts a partial record at the
+// end of the file.
+type Records struct {
+ file *os.File
+ size int
+ count int
+}
+
// Read returns record n, or false when the file does not reach it.
-func (r *records) Read(n int) ([]byte, bool, error) {
+func (r *Records) Read(n int) ([]byte, bool, error) {
if n >= r.count {
return nil, false, nil
}
}
// Write writes record n.
-func (r *records) Write(n int, record []byte) error {
+func (r *Records) Write(n int, record []byte) error {
_, err := r.file.WriteAt(record, int64(n*r.size))
if err != nil {
return fmt.Errorf("write index record %d: %w", n, err)
}
// Append writes the record after the last one.
-func (r *records) Append(record []byte) error {
+func (r *Records) Append(record []byte) error {
return r.Write(r.count, record)
}
// Close closes the file.
-func (r *records) Close() error {
+func (r *Records) Close() error {
return r.file.Close()
}
-
-// encodePosition writes place into the first positionBytes bytes of record.
-func encodePosition(record []byte, place position) {
- binary.LittleEndian.PutUint32(record[0:], place.day)
- binary.LittleEndian.PutUint32(record[4:], place.length)
- binary.LittleEndian.PutUint64(record[8:], place.offset)
-}
-
-// decodePosition reads a position from the first positionBytes bytes of record.
-func decodePosition(record []byte) position {
- return position{
- day: binary.LittleEndian.Uint32(record[0:]),
- length: binary.LittleEndian.Uint32(record[4:]),
- offset: binary.LittleEndian.Uint64(record[8:]),
- }
-}
"github.com/stretchr/testify/require"
)
-// syntheticMessages is the length of the synthetic log.
-const syntheticMessages = 100000
+// replayMessages is the length of the synthetic log.
+const replayMessages = 100000
// TestReplaySyntheticLog runs a synthetic log through the memory core with the
// fake summarizer. After each message it waits for the builds, as a turn does.
batches := 0
low := viewChatHigh
- long := make([]bool, syntheticMessages)
- for id := range syntheticMessages {
- if id == syntheticMessages/2 {
- memory = requireRestart(t, memory, root, clock)
+ long := make([]bool, replayMessages)
+ for id := range replayMessages {
+ if id == replayMessages/2 {
+ memory = replayRequireRestart(t, memory, root, clock)
}
clock.now = clock.now.Add(time.Minute)
- kind, text := syntheticMessage(random, id)
+ kind, text := replayMessage(random, id)
long[id] = len(kind)+len(": ")+len(text) > nodeBytes
- lines := viewLines(memory)
+ lines := replayViewLines(memory)
_, err = memory.Append(kind, text)
require.NoError(t, err)
- size, merged := requireViews(t, memory, lines, long, id%1000 == 0)
+ size, merged := replayRequireViews(t, memory, lines, long, id%1000 == 0)
if merged {
batches++
low = min(low, size)
require.NoError(t, memory.Wait(t.Context()))
}
- memory = requireRestart(t, memory, root, clock)
- require.Equal(t, syntheticMessages, memory.log.Len())
+ memory = replayRequireRestart(t, memory, root, clock)
+ require.Equal(t, replayMessages, memory.log.Len())
require.Greater(t, batches, 0)
require.NoError(t, memory.Wait(t.Context()))
- t.Logf("%d messages, %d nodes, %d batches in %s", syntheticMessages, memory.tree.Count(), batches, time.Since(start).Round(time.Second))
+ t.Logf("%d messages, %d nodes, %d batches in %s", replayMessages, memory.tree.Count(), batches, time.Since(start).Round(time.Second))
t.Logf("chat view: %d lines, %d bytes, lowest after a batch %d bytes", len(memory.chat.lines), memory.chat.Size(), low)
t.Logf("compaction view: %d lines, %d bytes", len(memory.cview.lines), memory.cview.Size())
t.Logf("heap in use: %d KB", stats.HeapInuse/1024)
}
-// syntheticMessage returns the kind and the text of message id. Most messages
-// are tool calls and their output. Users write short lines. One output in a
-// hundred is over 30 KB.
-func syntheticMessage(random *rand.Rand, id int) (string, string) {
+// replayRequireRestart closes memory, opens it again, and checks that the views
+// and view.json stay the same byte for byte.
+func replayRequireRestart(t *testing.T, memory *Memory, root string, clock *TestClock) *Memory {
+ t.Helper()
+
+ require.NoError(t, memory.Wait(t.Context()))
+
+ view := memory.View()
+ context := memory.Context(0, memory.log.Len())
+ saved, err := os.ReadFile(filepath.Join(root, "view.json"))
+ require.NoError(t, err)
+ require.NoError(t, memory.Close())
+
+ memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
+ require.NoError(t, err)
+ require.NoError(t, memory.Wait(t.Context()))
+ require.Equal(t, view, memory.View())
+ require.Equal(t, context, memory.Context(0, memory.log.Len()))
+
+ loaded, err := os.ReadFile(filepath.Join(root, "view.json"))
+ require.NoError(t, err)
+ require.Equal(t, saved, loaded)
+
+ return memory
+}
+
+// replayMessage returns the kind and the text of message id. Most messages are
+// tool calls and their output. Users write short lines. One output in a hundred
+// is over 30 KB.
+func replayMessage(random *rand.Rand, id int) (string, string) {
roll := random.IntN(100)
switch {
case roll < 5:
- return "user", syntheticText(id, 10+random.IntN(200))
+ return "user", replayText(id, 10+random.IntN(200))
case roll < 10:
- return "reply", syntheticText(id, 100+random.IntN(1500))
+ return "reply", replayText(id, 100+random.IntN(1500))
case roll < 12:
- return "note", syntheticText(id, 50+random.IntN(600))
+ return "note", replayText(id, 50+random.IntN(600))
case roll < 13:
- return "work", syntheticText(id, 300+random.IntN(3000))
+ return "work", replayText(id, 300+random.IntN(3000))
case roll < 56:
- return "tool", syntheticText(id, 60+random.IntN(400))
+ return "tool", replayText(id, 60+random.IntN(400))
case roll < 57:
- return "echo", syntheticText(id, 30000+random.IntN(30000))
+ return "echo", replayText(id, 30000+random.IntN(30000))
default:
- return "echo", syntheticText(id, 20+random.IntN(4000))
+ return "echo", replayText(id, 20+random.IntN(4000))
}
}
-// syntheticText returns a text of about size bytes that starts with its message
-// ID and holds some non-ASCII characters.
-func syntheticText(id, size int) string {
- head := fmt.Sprintf("m%d ", id)
- words := "ação data line café "
+// replayViewLines returns the number of lines in the chat view of memory.
+func replayViewLines(memory *Memory) int {
+ memory.mutex.Lock()
+ defer memory.mutex.Unlock()
- return head + strings.Repeat(words, max(1, (size-len(head))/len(words)))
+ return len(memory.chat.lines)
}
-// requireViews checks the views after a new message, and with full compares
-// them with the tree. Each view is at most its high limit, and a batch merges
-// the chat view to its low limit. The leaf of the new message can be built
-// after the append, so the sizes count it as the pending mark, as the append
-// did. It returns the size of the chat view and if a batch merged it, from the
-// number of lines before the message.
-func requireViews(t *testing.T, memory *Memory, lines int, long []bool, full bool) (int, bool) {
+// replayRequireViews checks the views after a new message, and with full
+// compares them with the tree. Each view is at most its high limit, and a batch
+// merges the chat view to its low limit. The leaf of the new message can be
+// built after the append, so the sizes count it as the pending mark, as the
+// append did. It returns the size of the chat view and if a batch merged it,
+// from the number of lines before the message.
+func replayRequireViews(t *testing.T, memory *Memory, lines int, long []bool, full bool) (int, bool) {
t.Helper()
memory.mutex.Lock()
defer memory.mutex.Unlock()
- size := appendSize(memory.chat)
+ size := replayAppendSize(memory.chat)
require.LessOrEqual(t, size, viewChatHigh)
- require.LessOrEqual(t, appendSize(memory.cview), viewCompactionHigh)
+ require.LessOrEqual(t, replayAppendSize(memory.cview), viewCompactionHigh)
merged := len(memory.chat.lines) < lines
if merged {
require.LessOrEqual(t, size, viewChatLow)
}
- requireOnlySummaries(t, memory, long, full)
+ replayRequireOnlySummaries(t, memory, long, full)
return size, merged
}
-// requireOnlySummaries checks each line of the chat view as the model reads it.
-// A built line holds its node, a long message has a summary of the fake
-// summarizer and not its text, and an unbuilt line holds only the pending mark.
-// With full, it also checks that the view keeps the node of each built line.
-// The caller holds the lock of memory.
-func requireOnlySummaries(t *testing.T, memory *Memory, long []bool, full bool) {
+// replayText returns a text of about size bytes that starts with its message ID
+// and holds some non-ASCII characters.
+func replayText(id, size int) string {
+ head := fmt.Sprintf("m%d ", id)
+ words := "ação data line café "
+
+ return head + strings.Repeat(words, max(1, (size-len(head))/len(words)))
+}
+
+// replayAppendSize returns the size of view with its last line as the pending
+// mark.
+func replayAppendSize(view *View) int {
+ last := view.lines[len(view.lines)-1]
+ return view.Size() - view.lineSize(last) + len(viewPending)
+}
+
+// replayRequireOnlySummaries checks each line of the chat view as the model
+// reads it. A built line holds its node, a long message has a summary of the
+// fake summarizer and not its text, and an unbuilt line holds only the pending
+// mark. With full, it also checks that the view keeps the node of each built
+// line. The caller holds the lock of memory.
+func replayRequireOnlySummaries(t *testing.T, memory *Memory, long []bool, full bool) {
t.Helper()
rendered := strings.Split(memory.chat.Render(), "\n")
head := fmt.Sprintf("%d+%d|", line.index<<line.level, 1<<line.level)
node, built := memory.chat.nodes[line]
if full {
- require.Equal(t, isBuilt(t, memory.tree, line.level, line.index), built)
+ require.Equal(t, treeIsBuilt(t, memory.tree, line.level, line.index), built)
}
if !built {
}
}
}
-
-// appendSize returns the size of view with its last line as the pending mark.
-func appendSize(view *View) int {
- last := view.lines[len(view.lines)-1]
- return view.Size() - view.lineSize(last) + len(viewPending)
-}
-
-// viewLines returns the number of lines in the chat view of memory.
-func viewLines(memory *Memory) int {
- memory.mutex.Lock()
- defer memory.mutex.Unlock()
-
- return len(memory.chat.lines)
-}
-
-// requireRestart closes memory, opens it again, and checks that the views and
-// view.json stay the same byte for byte.
-func requireRestart(t *testing.T, memory *Memory, root string, clock *testClock) *Memory {
- t.Helper()
-
- require.NoError(t, memory.Wait(t.Context()))
-
- view := memory.View()
- context := memory.Context(0, memory.log.Len())
- saved, err := os.ReadFile(filepath.Join(root, "view.json"))
- require.NoError(t, err)
- require.NoError(t, memory.Close())
-
- memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
- require.NoError(t, err)
- require.NoError(t, memory.Wait(t.Context()))
- require.Equal(t, view, memory.View())
- require.Equal(t, context, memory.Context(0, memory.log.Len()))
-
- loaded, err := os.ReadFile(filepath.Join(root, "view.json"))
- require.NoError(t, err)
- require.Equal(t, saved, loaded)
-
- return memory
-}
--- /dev/null
+package main
+
+import (
+ "context"
+ "errors"
+ "fmt"
+ "slices"
+ "strings"
+ "sync"
+ "unicode/utf8"
+)
+
+// 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 (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
+}
+
+// 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)
+}
+
+// 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
+}
+
+func newFakeSummarizer() *FakeSummarizer {
+ return &FakeSummarizer{}
+}
"golang.org/x/sys/unix"
)
-// syncer makes the writes of a chat durable by group commit. Each write goes to
+// openSyncer returns the syncer of the file system that holds dir. It syncs
+// with syncfs, which writes back only that file system.
+func openSyncer(dir string) (*Syncer, error) {
+ file, err := os.Open(dir)
+ if err != nil {
+ return nil, fmt.Errorf("open chat directory: %w", err)
+ }
+
+ s := newSyncer(func() error {
+ return unix.Syncfs(int(file.Fd()))
+ })
+ s.file = file
+
+ return s, nil
+}
+
+// Syncer makes the writes of a chat durable by group commit. Each write goes to
// the operating system at once, so a crash of the process loses nothing. After
// each write, Written starts a background sync of the file system of the chat,
// unless one runs already. Then the next sync starts when the running one ends.
// One sync makes every write before its start durable, so a sync covers many
-// writes, and there is never more than one sync at a time. A power loss loses at
-// most the writes after the start of the last finished sync.
+// writes, and there is never more than one sync at a time. A power loss loses
+// at most the writes after the start of the last finished sync.
//
// A failed sync can mean that the kernel dropped written data, so the error
// stays, and each later Flush returns it.
-type syncer struct {
+type Syncer struct {
mutex sync.Mutex
sync func() error
file *os.File
workers sync.WaitGroup
}
-// openSyncer returns the syncer of the file system that holds dir. It syncs
-// with syncfs, which writes back only that file system.
-func openSyncer(dir string) (*syncer, error) {
- file, err := os.Open(dir)
- if err != nil {
- return nil, fmt.Errorf("open chat directory: %w", err)
- }
-
- s := newSyncer(func() error {
- return unix.Syncfs(int(file.Fd()))
- })
- s.file = file
-
- return s, nil
-}
-
-// newSyncer returns a syncer that makes the writes durable with sync.
-func newSyncer(sync func() error) *syncer {
- return &syncer{
- sync: sync,
- done: make(chan struct{}),
- }
-}
-
// Written tells the syncer that a write went to the operating system, and
// starts a sync when none runs.
-func (s *syncer) Written() {
+func (s *Syncer) Written() {
s.mutex.Lock()
defer s.mutex.Unlock()
s.start()
}
-// start starts a sync of the writes so far when none runs and a write is not
-// durable. The caller holds the mutex.
-func (s *syncer) start() {
- if s.running || s.durable == s.written || s.failure != nil {
- return
- }
-
- s.running = true
- s.workers.Add(1)
- go s.run(s.written)
-}
-
-// run syncs once, so the writes up to target become durable, and starts the
-// next sync when writes came during this one.
-func (s *syncer) run(target uint64) {
- defer s.workers.Done()
-
- err := s.sync()
-
- s.mutex.Lock()
- defer s.mutex.Unlock()
-
- s.running = false
- if err != nil {
- s.failure = fmt.Errorf("sync chat files: %w", err)
- } else {
- s.durable = target
- }
-
- close(s.done)
- s.done = make(chan struct{})
- s.start()
-}
-
// Flush returns when every write before the call is durable.
-func (s *syncer) Flush(ctx context.Context) error {
+func (s *Syncer) Flush(ctx context.Context) error {
s.mutex.Lock()
target := s.written
s.mutex.Unlock()
// Close makes every write durable, waits for the syncs, and closes the
// directory.
-func (s *syncer) Close() error {
+func (s *Syncer) Close() error {
err := s.Flush(context.Background())
s.workers.Wait()
if s.file != nil {
return err
}
+
+// start starts a sync of the writes so far when none runs and a write is not
+// durable. The caller holds the mutex.
+func (s *Syncer) start() {
+ if s.running || s.durable == s.written || s.failure != nil {
+ return
+ }
+
+ s.running = true
+ s.workers.Add(1)
+ go s.run(s.written)
+}
+
+// run syncs once, so the writes up to target become durable, and starts the
+// next sync when writes came during this one.
+func (s *Syncer) run(target uint64) {
+ defer s.workers.Done()
+
+ err := s.sync()
+
+ s.mutex.Lock()
+ defer s.mutex.Unlock()
+
+ s.running = false
+ if err != nil {
+ s.failure = fmt.Errorf("sync chat files: %w", err)
+ } else {
+ s.durable = target
+ }
+
+ close(s.done)
+ s.done = make(chan struct{})
+ s.start()
+}
+
+// newSyncer returns a syncer that makes the writes durable with sync.
+func newSyncer(sync func() error) *Syncer {
+ return &Syncer{
+ sync: sync,
+ done: make(chan struct{}),
+ }
+}
"github.com/stretchr/testify/require"
)
-// childDirVariable names the chat directory of the writer child process of
-// TestKilledWriterLosesNoLoggedMessage.
-const childDirVariable = "ONECHAT_TEST_CHILD_DIR"
+// syncerChildDirVariable names the chat directory of the writer child process
+// of TestKilledWriterLosesNoLoggedMessage.
+const syncerChildDirVariable = "ONECHAT_TEST_CHILD_DIR"
func TestFlushWaitsForASyncThatStartsAfterTheWrite(t *testing.T) {
- gate := newGatedSync()
+ gate := newSyncerGatedSync()
s := newSyncer(gate.Sync)
s.Written()
gate.WaitStarted(t, 1)
}
func TestOneSyncCoversTheWritesDuringTheSyncBeforeIt(t *testing.T) {
- gate := newGatedSync()
+ gate := newSyncerGatedSync()
s := newSyncer(gate.Sync)
s.Written()
gate.WaitStarted(t, 1)
}
func TestFlushWithoutWritesNeedsNoSync(t *testing.T) {
- gate := newGatedSync()
+ gate := newSyncerGatedSync()
s := newSyncer(gate.Sync)
require.NoError(t, s.Flush(t.Context()))
require.NoError(t, s.Close())
func TestKilledWriterLosesNoLoggedMessage(t *testing.T) {
root := t.TempDir()
for round := range 3 {
- acknowledged := runKilledWriter(t, root, 200+100*round)
+ acknowledged := syncerRunKilledWriter(t, root, 200+100*round)
memory, err := openMemory(t.Context(), root, time.Now, newFakeSummarizer())
require.NoError(t, err)
for id := range memory.log.Len() {
message, err := memory.log.Message(id)
require.NoError(t, err)
- require.Equal(t, writerKind(id), message.Kind)
- require.Equal(t, writerText(id), message.Text)
+ require.Equal(t, syncerWriterKind(id), message.Kind)
+ require.Equal(t, syncerWriterText(id), message.Text)
}
require.NoError(t, memory.Close())
}
// TestWriterChild is the child process of TestKilledWriterLosesNoLoggedMessage.
-// It appends messages to the chat in the directory of childDirVariable and
-// prints the ID of each message after its Append returns, until it is killed.
+// It appends messages to the chat in the directory of syncerChildDirVariable
+// and prints the ID of each message after its Append returns, until it is
+// killed.
func TestWriterChild(t *testing.T) {
- root := os.Getenv(childDirVariable)
+ root := os.Getenv(syncerChildDirVariable)
if root == "" {
t.Skip("runs only as a child process")
}
require.NoError(t, err)
for id := memory.log.Len(); ; id++ {
- _, err = memory.Append(writerKind(id), writerText(id))
+ _, err = memory.Append(syncerWriterKind(id), syncerWriterText(id))
require.NoError(t, err)
_, err = fmt.Println(id)
}
}
-// gatedSync is a sync function that blocks until the test releases it.
-type gatedSync struct {
+// SyncerGatedSync is a sync function that blocks until the test releases it.
+type SyncerGatedSync struct {
mutex sync.Mutex
count int
started chan struct{}
release chan struct{}
}
-func newGatedSync() *gatedSync {
- return &gatedSync{
- started: make(chan struct{}, 100),
- release: make(chan struct{}),
- }
-}
-
-func (g *gatedSync) Sync() error {
+func (g *SyncerGatedSync) Sync() error {
g.mutex.Lock()
g.count++
g.mutex.Unlock()
}
// WaitStarted waits until sync number n has started.
-func (g *gatedSync) WaitStarted(t *testing.T, n int) {
+func (g *SyncerGatedSync) WaitStarted(t *testing.T, n int) {
t.Helper()
require.Eventually(t, func() bool {
}
// Release lets the running sync end.
-func (g *gatedSync) Release() {
+func (g *SyncerGatedSync) Release() {
<-g.started
g.release <- struct{}{}
}
-func (g *gatedSync) Count() int {
+func (g *SyncerGatedSync) Count() int {
g.mutex.Lock()
defer g.mutex.Unlock()
return g.count
}
-// runKilledWriter starts the writer child on the chat in root, kills it with
-// SIGKILL after it acknowledged count messages, and returns the number of
+func newSyncerGatedSync() *SyncerGatedSync {
+ return &SyncerGatedSync{
+ started: make(chan struct{}, 100),
+ release: make(chan struct{}),
+ }
+}
+
+// syncerRunKilledWriter starts the writer child on the chat in root, kills it
+// with SIGKILL after it acknowledged count messages, and returns the number of
// messages that it acknowledged.
-func runKilledWriter(t *testing.T, root string, count int) int {
+func syncerRunKilledWriter(t *testing.T, root string, count int) int {
t.Helper()
command := exec.Command(os.Args[0], "-test.run=^TestWriterChild$")
- command.Env = append(os.Environ(), childDirVariable+"="+root)
+ command.Env = append(os.Environ(), syncerChildDirVariable+"="+root)
stdout, err := command.StdoutPipe()
require.NoError(t, err)
require.NoError(t, command.Start())
return acknowledged
}
-// writerKind returns the kind of message id of the writer child.
-func writerKind(id int) string {
+// syncerWriterKind returns the kind of message id of the writer child.
+func syncerWriterKind(id int) string {
return []string{"user", "tool", "echo", "reply"}[id%4]
}
-// writerText returns the text of message id of the writer child. Some texts
-// are large, so a kill can come during the write of one line.
-func writerText(id int) string {
+// syncerWriterText returns the text of message id of the writer child. Some
+// texts are large, so a kill can come during the write of one line.
+func syncerWriterText(id int) string {
size := 40 + id%7*30
if id%5 == 4 {
size = 64*1024 + id
// few bytes over it, because the view measures the real sizes.
const nodeBytes = 512
-// Node is one line of the summary tree. Node (l, i) covers the 2^l messages
-// from i*2^l on. A node of level 0 holds one message. A node of a higher level
-// merges its two children, (l-1, 2i) and (l-1, 2i+1). Size is the UTF-8 byte
-// length of Text.
-type Node struct {
- Level int `json:"l"`
- Index int `json:"i"`
- Text string `json:"text"`
- Size int `json:"size"`
+// treeRecordBytes is the size of a record of index/tree.idx: the position of
+// the node line, the node index, and the node level.
+const treeRecordBytes = 32
+
+// openTree opens the tree of the chat in root and indexes its new lines. A chat
+// of total messages cannot have a node that covers a later message.
+func openTree(ctx context.Context, root string, total int, clock func() time.Time) (*Tree, error) {
+ for _, dir := range []string{"tree", "index"} {
+ err := os.MkdirAll(filepath.Join(root, dir), 0o700)
+ if err != nil {
+ return nil, fmt.Errorf("create %s directory: %w", dir, err)
+ }
+ }
+
+ t := &Tree{
+ clock: clock,
+ journal: Journal{
+ dir: filepath.Join(root, "tree"),
+ },
+ dir: filepath.Join(root, "index"),
+ levels: make(map[int]*Records),
+ }
+
+ err := t.open(ctx, total)
+ if err != nil {
+ return nil, errors.Join(fmt.Errorf("index tree: %w", err), t.Close())
+ }
+
+ return t, nil
}
-// nodeKey names the node (l, i).
-type nodeKey struct {
- level int
- index int
+// leafSource returns the text that the leaf of message holds or summarizes.
+func leafSource(message Message) string {
+ return message.Kind + ": " + message.Text
}
-// treeRecordBytes is the size of a record of index/tree.idx: the position of
-// the node line, the node index, and the node level.
-const treeRecordBytes = 32
+// freeLeaf returns the leaf of message when its source fits in a node. Such a
+// leaf holds the source word for word and needs no model call.
+func freeLeaf(message Message) (Node, bool) {
+ text := leafSource(message)
+ if len(text) > nodeBytes {
+ return Node{}, false
+ }
+
+ return Node{
+ Level: 0,
+ Index: message.ID,
+ Text: text,
+ Size: len(text),
+ }, true
+}
+
+// freeMerge returns the parent of two sibling nodes when both texts, joined by
+// a newline, fit in a node. Such a parent needs no model call.
+func freeMerge(left, right Node) (Node, bool) {
+ if left.Size+1+right.Size > nodeBytes {
+ return Node{}, false
+ }
+
+ text := left.Text + "\n" + right.Text
+
+ return Node{
+ Level: left.Level + 1,
+ Index: left.Index / 2,
+ Text: text,
+ Size: len(text),
+ }, true
+}
+
+// nodeEnd returns the ID after the last message of node (l, i).
+func nodeEnd(level, index int) int {
+ return (index + 1) << level
+}
// Tree holds the built nodes of one chat. It writes one JSON line for each node
// to tree/YYYY-MM-DD.jsonl, named for the local day of the build. Each node is
// built again from the day files.
type Tree struct {
clock func() time.Time
- journal journal
+ journal Journal
dir string
- order *records
- levels map[int]*records
+ order *Records
+ levels map[int]*Records
failure error
}
-// openTree opens the tree of the chat in root and indexes its new lines. A chat
-// of total messages cannot have a node that covers a later message.
-func openTree(ctx context.Context, root string, total int, clock func() time.Time) (*Tree, error) {
- for _, dir := range []string{"tree", "index"} {
- err := os.MkdirAll(filepath.Join(root, dir), 0o700)
- if err != nil {
- return nil, fmt.Errorf("create %s directory: %w", dir, err)
- }
+// Close closes the files of the tree.
+func (t *Tree) Close() error {
+ err := t.journal.Close()
+ if t.order != nil {
+ err = errors.Join(err, t.order.Close())
+ t.order = nil
}
- t := &Tree{
- clock: clock,
- journal: journal{
- dir: filepath.Join(root, "tree"),
- },
- dir: filepath.Join(root, "index"),
- levels: make(map[int]*records),
+ for level, index := range t.levels {
+ err = errors.Join(err, index.Close())
+ delete(t.levels, level)
}
- err := t.open(ctx, total)
+ return err
+}
+
+// Add writes a new node to the tree.
+func (t *Tree) Add(node Node) error {
+ if t.failure != nil {
+ return fmt.Errorf("tree needs a reopen after a failed write: %w", t.failure)
+ }
+
+ err := t.check(node)
if err != nil {
- return nil, errors.Join(fmt.Errorf("index tree: %w", err), t.Close())
+ return err
}
- return t, nil
+ line, err := json.Marshal(node)
+ if err != nil {
+ return fmt.Errorf("encode node: %w", err)
+ }
+
+ place, err := t.journal.Write(t.clock().Format(time.DateOnly), append(line, '\n'))
+ if err == nil {
+ err = t.addIndex(NodeKey{node.Level, node.Index}, place)
+ }
+
+ if err != nil {
+ t.failure = err
+ return fmt.Errorf("write node (%d, %d): %w", node.Level, node.Index, err)
+ }
+
+ return nil
+}
+
+// Built tells if the node (l, i) is built.
+func (t *Tree) Built(level, index int) (bool, error) {
+ _, built, err := t.place(level, index)
+ return built, err
+}
+
+// Node reads the node (l, i) if it is built.
+func (t *Tree) Node(level, index int) (Node, bool, error) {
+ place, built, err := t.place(level, index)
+ if err != nil || !built {
+ return Node{}, false, err
+ }
+
+ line, err := t.journal.Read(place)
+ if err != nil {
+ return Node{}, false, fmt.Errorf("read node (%d, %d): %w", level, index, err)
+ }
+
+ var node Node
+ err = json.Unmarshal(line, &node)
+ if err != nil {
+ return Node{}, false, fmt.Errorf("decode node (%d, %d): %w", level, index, err)
+ }
+
+ if node.Level != level || node.Index != index {
+ return Node{}, false, fmt.Errorf("index of node (%d, %d) points to node (%d, %d)", level, index, node.Level, node.Index)
+ }
+
+ return node, true, nil
+}
+
+// Keys returns the keys of the built nodes, ordered by level and then by index.
+func (t *Tree) Keys() ([]NodeKey, error) {
+ keys := make([]NodeKey, 0, t.order.count)
+ for n := range t.order.count {
+ record, _, err := t.order.Read(n)
+ if err != nil {
+ return nil, err
+ }
+
+ key, _ := treeDecodeRecord(record)
+ keys = append(keys, key)
+ }
+
+ slices.SortFunc(keys, func(a, b NodeKey) int {
+ return cmp.Or(cmp.Compare(a.level, b.level), cmp.Compare(a.index, b.index))
+ })
+
+ return keys, nil
}
func (t *Tree) open(ctx context.Context, total int) error {
return err
}
- start := position{}
+ start := Position{}
if t.order.count == 0 {
err = t.removeLevels()
if err != nil {
// A crash can stop an add between the two index writes, so the
// last node gets its level record again.
- node, last := decodeTreeRecord(record)
+ node, last := treeDecodeRecord(record)
err = t.writeLevel(node, last)
if err != nil {
return err
}
- start = position{
+ start = Position{
day: last.day,
offset: last.offset + uint64(last.length),
}
}
- return t.journal.Scan(ctx, start, func(line []byte, place position) error {
+ return t.journal.Scan(ctx, start, func(line []byte, place Position) error {
var node Node
err := json.Unmarshal(line, &node)
if err != nil {
return fmt.Errorf("node (%d, %d) covers messages after the log", node.Level, node.Index)
}
- return t.addIndex(nodeKey{node.Level, node.Index}, place)
+ return t.addIndex(NodeKey{node.Level, node.Index}, place)
})
}
-// removeLevels removes the level indexes, which are not valid without the order
-// index.
-func (t *Tree) removeLevels() error {
- paths, err := filepath.Glob(filepath.Join(t.dir, "tree-*.idx"))
- if err != nil {
- return fmt.Errorf("list level indexes: %w", err)
- }
-
- for _, path := range paths {
- err = os.Remove(path)
- if err != nil {
- return fmt.Errorf("remove level index: %w", err)
- }
- }
-
- return nil
-}
-
-// writeLevel writes the position of node to its level index.
-func (t *Tree) writeLevel(node nodeKey, place position) error {
- level, err := t.level(node.level)
- if err != nil {
- return err
- }
-
- record := make([]byte, positionBytes)
- encodePosition(record, place)
-
- return level.Write(node.index, record)
-}
-
-// level returns the index of level, and opens it at first use.
-func (t *Tree) level(level int) (*records, error) {
- index, open := t.levels[level]
- if open {
- return index, nil
- }
-
- index, err := openRecords(filepath.Join(t.dir, fmt.Sprintf("tree-%d.idx", level)), positionBytes)
- if err != nil {
- return nil, err
- }
-
- t.levels[level] = index
-
- return index, nil
-}
-
// check makes sure that node is valid and can join the tree now.
func (t *Tree) check(node Node) error {
if node.Level < 0 || node.Level >= bits.UintSize-2 || node.Index < 0 {
}
// addIndex adds the node line at place to both indexes.
-func (t *Tree) addIndex(node nodeKey, place position) error {
+func (t *Tree) addIndex(node NodeKey, place Position) error {
record := make([]byte, treeRecordBytes)
encodePosition(record, place)
binary.LittleEndian.PutUint64(record[positionBytes:], uint64(node.index))
return t.writeLevel(node, place)
}
-// Close closes the files of the tree.
-func (t *Tree) Close() error {
- err := t.journal.Close()
- if t.order != nil {
- err = errors.Join(err, t.order.Close())
- t.order = nil
- }
-
- for level, index := range t.levels {
- err = errors.Join(err, index.Close())
- delete(t.levels, level)
- }
-
- return err
-}
-
-// Add writes a new node to the tree.
-func (t *Tree) Add(node Node) error {
- if t.failure != nil {
- return fmt.Errorf("tree needs a reopen after a failed write: %w", t.failure)
- }
-
- err := t.check(node)
- if err != nil {
- return err
- }
-
- line, err := json.Marshal(node)
- if err != nil {
- return fmt.Errorf("encode node: %w", err)
- }
-
- place, err := t.journal.Write(t.clock().Format(time.DateOnly), append(line, '\n'))
- if err == nil {
- err = t.addIndex(nodeKey{node.Level, node.Index}, place)
- }
-
- if err != nil {
- t.failure = err
- return fmt.Errorf("write node (%d, %d): %w", node.Level, node.Index, err)
- }
-
- return nil
-}
-
-// Built tells if the node (l, i) is built.
-func (t *Tree) Built(level, index int) (bool, error) {
- _, built, err := t.place(level, index)
- return built, err
-}
-
// place returns the position of the node (l, i) if it is built.
-func (t *Tree) place(level, index int) (position, bool, error) {
+func (t *Tree) place(level, index int) (Position, bool, error) {
if level < 0 || level >= bits.UintSize-2 || index < 0 {
- return position{}, false, nil
+ return Position{}, false, nil
}
records, err := t.level(level)
if err != nil {
- return position{}, false, err
+ return Position{}, false, err
}
record, found, err := records.Read(index)
if err != nil || !found {
- return position{}, false, err
+ return Position{}, false, err
}
place := decodePosition(record)
return place, place.length > 0, nil
}
-// Node reads the node (l, i) if it is built.
-func (t *Tree) Node(level, index int) (Node, bool, error) {
- place, built, err := t.place(level, index)
- if err != nil || !built {
- return Node{}, false, err
- }
-
- line, err := t.journal.Read(place)
+// removeLevels removes the level indexes, which are not valid without the order
+// index.
+func (t *Tree) removeLevels() error {
+ paths, err := filepath.Glob(filepath.Join(t.dir, "tree-*.idx"))
if err != nil {
- return Node{}, false, fmt.Errorf("read node (%d, %d): %w", level, index, err)
+ return fmt.Errorf("list level indexes: %w", err)
}
- var node Node
- err = json.Unmarshal(line, &node)
- if err != nil {
- return Node{}, false, fmt.Errorf("decode node (%d, %d): %w", level, index, err)
+ for _, path := range paths {
+ err = os.Remove(path)
+ if err != nil {
+ return fmt.Errorf("remove level index: %w", err)
+ }
}
- if node.Level != level || node.Index != index {
- return Node{}, false, fmt.Errorf("index of node (%d, %d) points to node (%d, %d)", level, index, node.Level, node.Index)
+ return nil
+}
+
+// writeLevel writes the position of node to its level index.
+func (t *Tree) writeLevel(node NodeKey, place Position) error {
+ level, err := t.level(node.level)
+ if err != nil {
+ return err
}
- return node, true, nil
-}
+ record := make([]byte, positionBytes)
+ encodePosition(record, place)
-// Count returns the number of built nodes.
-func (t *Tree) Count() int {
- return t.order.count
+ return level.Write(node.index, record)
}
-// Keys returns the keys of the built nodes, ordered by level and then by index.
-func (t *Tree) Keys() ([]nodeKey, error) {
- keys := make([]nodeKey, 0, t.order.count)
- for n := range t.order.count {
- record, _, err := t.order.Read(n)
- if err != nil {
- return nil, err
- }
+// level returns the index of level, and opens it at first use.
+func (t *Tree) level(level int) (*Records, error) {
+ index, open := t.levels[level]
+ if open {
+ return index, nil
+ }
- key, _ := decodeTreeRecord(record)
- keys = append(keys, key)
+ index, err := openRecords(filepath.Join(t.dir, fmt.Sprintf("tree-%d.idx", level)), positionBytes)
+ if err != nil {
+ return nil, err
}
- slices.SortFunc(keys, func(a, b nodeKey) int {
- return cmp.Or(cmp.Compare(a.level, b.level), cmp.Compare(a.index, b.index))
- })
+ t.levels[level] = index
- return keys, nil
+ return index, nil
}
-// leafSource returns the text that the leaf of message holds or summarizes.
-func leafSource(message Message) string {
- return message.Kind + ": " + message.Text
+// Count returns the number of built nodes.
+func (t *Tree) Count() int {
+ return t.order.count
}
-// freeLeaf returns the leaf of message when its source fits in a node. Such a
-// leaf holds the source word for word and needs no model call.
-func freeLeaf(message Message) (Node, bool) {
- text := leafSource(message)
- if len(text) > nodeBytes {
- return Node{}, false
- }
-
- return Node{
- Level: 0,
- Index: message.ID,
- Text: text,
- Size: len(text),
- }, true
+// Node is one line of the summary tree. Node (l, i) covers the 2^l messages
+// from i*2^l on. A node of level 0 holds one message. A node of a higher level
+// merges its two children, (l-1, 2i) and (l-1, 2i+1). Size is the UTF-8 byte
+// length of Text.
+type Node struct {
+ Level int `json:"l"`
+ Index int `json:"i"`
+ Text string `json:"text"`
+ Size int `json:"size"`
}
-// freeMerge returns the parent of two sibling nodes when both texts, joined by
-// a newline, fit in a node. Such a parent needs no model call.
-func freeMerge(left, right Node) (Node, bool) {
- if left.Size+1+right.Size > nodeBytes {
- return Node{}, false
- }
-
- text := left.Text + "\n" + right.Text
-
- return Node{
- Level: left.Level + 1,
- Index: left.Index / 2,
- Text: text,
- Size: len(text),
- }, true
+// NodeKey names the node (l, i).
+type NodeKey struct {
+ level int
+ index int
}
-// decodeTreeRecord reads the key and the position of a record of
+// treeDecodeRecord reads the key and the position of a record of
// index/tree.idx.
-func decodeTreeRecord(record []byte) (nodeKey, position) {
- key := nodeKey{
+func treeDecodeRecord(record []byte) (NodeKey, Position) {
+ key := NodeKey{
level: int(binary.LittleEndian.Uint32(record[positionBytes+8:])),
index: int(binary.LittleEndian.Uint64(record[positionBytes:])),
}
return key, decodePosition(record)
}
-
-// nodeEnd returns the ID after the last message of node (l, i).
-func nodeEnd(level, index int) int {
- return (index + 1) << level
-}
}
func TestFreeMergeJoinsShortSiblings(t *testing.T) {
- left := testNode(1, 2, "left")
- right := testNode(1, 3, "right")
+ left := treeTestNode(1, 2, "left")
+ right := treeTestNode(1, 3, "right")
node, free := freeMerge(left, right)
require.True(t, free)
- require.Equal(t, testNode(2, 1, "left\nright"), node)
+ require.Equal(t, treeTestNode(2, 1, "left\nright"), node)
- right = testNode(1, 3, strings.Repeat("r", nodeBytes-len("left")))
+ right = treeTestNode(1, 3, strings.Repeat("r", nodeBytes-len("left")))
_, free = freeMerge(left, right)
require.False(t, free)
}
require.NoError(t, tree.Close())
}()
- require.NoError(t, tree.Add(testNode(0, 0, "user: a")))
- require.ErrorContains(t, tree.Add(testNode(0, 0, "user: b")), "already built")
- require.ErrorContains(t, tree.Add(testNode(1, 0, "a b")), "child that is not built")
+ require.NoError(t, tree.Add(treeTestNode(0, 0, "user: a")))
+ require.ErrorContains(t, tree.Add(treeTestNode(0, 0, "user: b")), "already built")
+ require.ErrorContains(t, tree.Add(treeTestNode(1, 0, "a b")), "child that is not built")
- require.NoError(t, tree.Add(testNode(0, 1, "reply: b")))
- require.NoError(t, tree.Add(testNode(1, 0, "a b")))
+ require.NoError(t, tree.Add(treeTestNode(0, 1, "reply: b")))
+ require.NoError(t, tree.Add(treeTestNode(1, 0, "a b")))
- require.Equal(t, "user: a", readNode(t, tree, 0, 0).Text)
+ require.Equal(t, "user: a", treeReadNode(t, tree, 0, 0).Text)
}
func TestTreeReopenLoadsTheSameNodes(t *testing.T) {
require.NoError(t, err)
nodes := []Node{
- testNode(0, 1, "reply: b"),
- testNode(0, 0, "user: a"),
- testNode(1, 0, "user: a\nreply: b"),
+ treeTestNode(0, 1, "reply: b"),
+ treeTestNode(0, 0, "user: a"),
+ treeTestNode(1, 0, "user: a\nreply: b"),
}
for _, node := range nodes {
require.NoError(t, tree.Add(node))
require.Equal(t, len(nodes), tree.Count())
for _, node := range nodes {
- require.Equal(t, node, readNode(t, tree, node.Level, node.Index))
+ require.Equal(t, node, treeReadNode(t, tree, node.Level, node.Index))
}
}
require.NoError(t, err)
nodes := []Node{
- testNode(0, 1, "reply: b"),
- testNode(0, 0, "user: a"),
- testNode(1, 0, "user: a\nreply: b"),
+ treeTestNode(0, 1, "reply: b"),
+ treeTestNode(0, 0, "user: a"),
+ treeTestNode(1, 0, "user: a\nreply: b"),
}
for _, node := range nodes {
require.NoError(t, tree.Add(node))
require.Equal(t, len(nodes), tree.Count())
for _, node := range nodes {
- require.Equal(t, node, readNode(t, tree, node.Level, node.Index))
+ require.Equal(t, node, treeReadNode(t, tree, node.Level, node.Index))
}
}
tree, err := openTree(t.Context(), root, 2, clock.Now)
require.NoError(t, err)
- require.NoError(t, tree.Add(testNode(0, 0, "user: a")))
- require.NoError(t, tree.Add(testNode(0, 1, "reply: b")))
+ require.NoError(t, tree.Add(treeTestNode(0, 0, "user: a")))
+ require.NoError(t, tree.Add(treeTestNode(0, 1, "reply: b")))
require.NoError(t, tree.Close())
// A crash after the line write and in the middle of the index writes.
}()
require.Equal(t, 2, tree.Count())
- require.Equal(t, "reply: b", readNode(t, tree, 0, 1).Text)
- require.NoError(t, tree.Add(testNode(1, 0, "user: a\nreply: b")))
+ require.Equal(t, "reply: b", treeReadNode(t, tree, 0, 1).Text)
+ require.NoError(t, tree.Add(treeTestNode(1, 0, "user: a\nreply: b")))
}
-// isBuilt tells if the node (l, i) of tree is built.
-func isBuilt(t *testing.T, tree *Tree, level, index int) bool {
- t.Helper()
-
- built, err := tree.Built(level, index)
- require.NoError(t, err)
-
- return built
+func treeTestNode(level, index int, text string) Node {
+ return Node{
+ Level: level,
+ Index: index,
+ Text: text,
+ Size: len(text),
+ }
}
-// readNode reads the node (l, i) of tree, which must be built.
-func readNode(t *testing.T, tree *Tree, level, index int) Node {
+// treeReadNode reads the node (l, i) of tree, which must be built.
+func treeReadNode(t *testing.T, tree *Tree, level, index int) Node {
t.Helper()
node, built, err := tree.Node(level, index)
return node
}
-func testNode(level, index int, text string) Node {
- return Node{
- Level: level,
- Index: index,
- Text: text,
- Size: len(text),
- }
+// treeIsBuilt tells if the node (l, i) of tree is built.
+func treeIsBuilt(t *testing.T, tree *Tree, level, index int) bool {
+ t.Helper()
+
+ built, err := tree.Built(level, index)
+ require.NoError(t, err)
+
+ return built
}
return &View{
high: high,
low: low,
- nodes: make(map[nodeKey]Node),
+ nodes: make(map[NodeKey]Node),
}
}
v := newView(high, low)
v.merging = state.Merging
for _, pair := range state.Lines {
- line := nodeKey{pair[0], pair[1]}
+ line := NodeKey{pair[0], pair[1]}
if line.level < 0 || line.level >= bits.UintSize-2 || line.index < 0 || line.index<<line.level != v.end {
return nil, fmt.Errorf("view line (%d, %d) does not start at message %d", line.level, line.index, v.end)
}
type View struct {
high int
low int
- lines []nodeKey
- nodes map[nodeKey]Node
+ lines []NodeKey
+ nodes map[NodeKey]Node
end int
merging bool
}
// Append adds the leaf of the next message as a line and merges when the view
// is too large. It tells if a merge changed the view.
func (v *View) Append(tree ViewNodeReader) (bool, error) {
- line := nodeKey{0, v.end}
+ line := NodeKey{0, v.end}
_, err := v.load(tree, line)
if err != nil {
return false, err
// either nest or do not overlap, so the result covers the messages without a
// gap.
func (v *View) Coarsen(other *View, tree ViewNodeReader) error {
- lines := make([]nodeKey, 0, len(v.lines))
- nodes := make(map[nodeKey]Node, len(v.nodes))
+ lines := make([]NodeKey, 0, len(v.lines))
+ nodes := make(map[NodeKey]Node, len(v.nodes))
mine, theirs := 0, 0
for start := 0; start < v.end; {
line := v.lines[mine]
// Refresh keeps node when it is the node of a pending line.
func (v *View) Refresh(node Node) {
- key := nodeKey{node.Level, node.Index}
+ key := NodeKey{node.Level, node.Index}
for n := len(v.lines) - 1; n >= 0; n-- {
line := v.lines[n]
if line == key {
}
// load keeps the node of line when it is built, and tells if it is.
-func (v *View) load(tree ViewNodeReader, line nodeKey) (bool, error) {
+func (v *View) load(tree ViewNodeReader, line NodeKey) (bool, error) {
node, built, err := tree.Node(line.level, line.index)
if err != nil || !built {
return false, err
}
left, right := v.lines[best], v.lines[best+1]
- parent := nodeKey{left.level + 1, left.index / 2}
+ parent := NodeKey{left.level + 1, left.index / 2}
_, err = v.load(tree, parent)
if err != nil {
return merged, err
}
// renderLine writes one view line, "id+n|text", and a newline.
-func (v *View) renderLine(text *strings.Builder, line nodeKey) {
+func (v *View) renderLine(text *strings.Builder, line NodeKey) {
body := viewPending
node, built := v.nodes[line]
if built {
}
// lineSize returns the size of a view line.
-func (v *View) lineSize(line nodeKey) int {
+func (v *View) lineSize(line NodeKey) int {
node, built := v.nodes[line]
if !built {
return len(viewPending)
func TestViewMergesOnlyBuiltParents(t *testing.T) {
tree := ViewPartTree{
- sizes: map[nodeKey]int{
+ sizes: map[NodeKey]int{
{0, 0}: 1,
{0, 1}: 1,
{0, 2}: 1,
require.NoError(t, err)
}
- require.Equal(t, []nodeKey{{0, 0}, {0, 1}, {0, 2}, {0, 3}}, view.lines)
+ require.Equal(t, []NodeKey{{0, 0}, {0, 1}, {0, 2}, {0, 3}}, view.lines)
require.True(t, view.merging)
- tree.sizes[nodeKey{1, 0}] = 1
+ tree.sizes[NodeKey{1, 0}] = 1
_, err := view.Append(tree)
require.NoError(t, err)
- require.Equal(t, []nodeKey{{1, 0}, {0, 2}, {0, 3}, {0, 4}}, view.lines)
+ require.Equal(t, []NodeKey{{1, 0}, {0, 2}, {0, 3}, {0, 4}}, view.lines)
require.True(t, view.merging)
- tree.sizes[nodeKey{1, 1}] = 1
- tree.sizes[nodeKey{1, 2}] = 1
- tree.sizes[nodeKey{2, 0}] = 1
+ tree.sizes[NodeKey{1, 1}] = 1
+ tree.sizes[NodeKey{1, 2}] = 1
+ tree.sizes[NodeKey{2, 0}] = 1
_, err = view.Append(tree)
require.NoError(t, err)
- require.Equal(t, []nodeKey{{2, 0}, {1, 2}}, view.lines)
+ require.Equal(t, []NodeKey{{2, 0}, {1, 2}}, view.lines)
require.False(t, view.merging)
}
func TestViewRendersPendingLinesWithoutTheirMessage(t *testing.T) {
tree := ViewPartTree{
- sizes: map[nodeKey]int{
+ sizes: map[NodeKey]int{
{0, 0}: 3,
{0, 1}: 3,
{1, 0}: 7,
require.NoError(t, err)
require.NoError(t, view.Coarsen(other, ViewBuiltTree{}))
- require.Equal(t, []nodeKey{{1, 0}, {1, 1}, {0, 4}, {0, 5}}, view.lines)
+ require.Equal(t, []NodeKey{{1, 0}, {1, 1}, {0, 4}, {0, 5}}, view.lines)
view.low = 2
require.NoError(t, view.Coarsen(other, ViewBuiltTree{}))
- require.Equal(t, []nodeKey{{2, 0}, {1, 2}}, view.lines)
+ require.Equal(t, []NodeKey{{2, 0}, {1, 2}}, view.lines)
require.False(t, view.merging)
}
func TestViewContextStopsAtTheFirstPendingLine(t *testing.T) {
tree := ViewPartTree{
- sizes: map[nodeKey]int{
+ sizes: map[NodeKey]int{
{0, 0}: 1,
{0, 1}: 1,
{1, 0}: 1,
}
func (ViewBuiltTree) Node(level, index int) (Node, bool, error) {
- return testNode(level, index, "a"), true, nil
+ return treeTestNode(level, index, "a"), true, nil
}
// ViewPushEntry is one entry of the rollback state list of the reference
// ViewPartTree is a tree with the built nodes of its sizes.
type ViewPartTree struct {
- sizes map[nodeKey]int
+ sizes map[NodeKey]int
}
func (p ViewPartTree) Built(level, index int) (bool, error) {
- _, built := p.sizes[nodeKey{level, index}]
+ _, built := p.sizes[NodeKey{level, index}]
return built, nil
}
func (p ViewPartTree) Node(level, index int) (Node, bool, error) {
- size, built := p.sizes[nodeKey{level, index}]
+ size, built := p.sizes[NodeKey{level, index}]
if !built {
return Node{}, false, nil
}
text = "aaa\naaa"
}
- return testNode(level, index, text), true, nil
+ return treeTestNode(level, index, text), true, nil
}
// viewPush is the push of rollback_state_list.js.
projects / onechat / commitdiff