]> git.giacomolli.uk - onechat/blob - journal.go
Say that the agent can change its instructions
[onechat] / journal.go
1 package main
2
3 import (
4 "bufio"
5 "context"
6 "errors"
7 "fmt"
8 "io"
9 "os"
10 "path/filepath"
11 "slices"
12 "strconv"
13 "strings"
14 )
15
16 // Position names one line of a day file: the day as YYYYMMDD, the byte offset
17 // of the line, and its length with the newline.
18 type Position struct {
19 day uint32
20 offset uint64
21 length uint32
22 }
23
24 // Journal appends JSON lines to one file for each day in dir, named
25 // YYYY-MM-DD.jsonl, and reads them back by their position. It keeps the file of
26 // the last day open for writes and each file that it read open for reads.
27 type Journal struct {
28 dir string
29 file *os.File
30 day uint32
31 size uint64
32 readers map[uint32]*os.File
33 }
34
35 // Write appends line to the file of day, YYYY-MM-DD, with one write to the
36 // operating system, without a buffer. A crash of the process then loses
37 // nothing. It returns the position of the line.
38 func (j *Journal) Write(day string, line []byte) (Position, error) {
39 number, err := journalDayNumber(day)
40 if err != nil {
41 return Position{}, err
42 }
43
44 if number != j.day {
45 err = j.closeWriter()
46 if err != nil {
47 return Position{}, err
48 }
49
50 j.file, err = os.OpenFile(filepath.Join(j.dir, day+".jsonl"), os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o600)
51 if err != nil {
52 return Position{}, fmt.Errorf("open day file: %w", err)
53 }
54
55 info, err := j.file.Stat()
56 if err != nil {
57 return Position{}, fmt.Errorf("read day file size: %w", err)
58 }
59
60 j.day = number
61 j.size = uint64(info.Size())
62 }
63
64 _, err = j.file.Write(line)
65 if err != nil {
66 return Position{}, fmt.Errorf("write line: %w", err)
67 }
68
69 place := Position{
70 day: number,
71 offset: j.size,
72 length: uint32(len(line)),
73 }
74 j.size += uint64(len(line))
75
76 return place, nil
77 }
78
79 // Read returns the line at place without its newline.
80 func (j *Journal) Read(place Position) ([]byte, error) {
81 reader, open := j.readers[place.day]
82 if !open {
83 var err error
84 reader, err = os.Open(filepath.Join(j.dir, journalDayName(place.day)+".jsonl"))
85 if err != nil {
86 return nil, fmt.Errorf("open day file: %w", err)
87 }
88
89 if j.readers == nil {
90 j.readers = make(map[uint32]*os.File)
91 }
92
93 j.readers[place.day] = reader
94 }
95
96 line := make([]byte, place.length)
97 _, err := reader.ReadAt(line, int64(place.offset))
98 if err != nil {
99 return nil, fmt.Errorf("read line of %s: %w", journalDayName(place.day), err)
100 }
101
102 if len(line) == 0 || line[len(line)-1] != '\n' {
103 return nil, fmt.Errorf("line of %s at %d does not end its record", journalDayName(place.day), place.offset)
104 }
105
106 return line[:len(line)-1], nil
107 }
108
109 // Close closes the open day files.
110 func (j *Journal) Close() error {
111 err := j.closeWriter()
112 for day, reader := range j.readers {
113 err = errors.Join(err, reader.Close())
114 delete(j.readers, day)
115 }
116
117 return err
118 }
119
120 // Scan gives each line of the day files from place on to add, with its
121 // position, oldest day first. The zero position starts at the first line. A
122 // crash during a write can leave a part of a line at the end of the last day
123 // file. Its write never returned, so Scan cuts it off.
124 func (j *Journal) Scan(ctx context.Context, place Position, add func([]byte, Position) error) error {
125 paths, err := filepath.Glob(filepath.Join(j.dir, "*.jsonl"))
126 if err != nil {
127 return fmt.Errorf("list day files: %w", err)
128 }
129
130 slices.Sort(paths)
131 for n, path := range paths {
132 day, err := journalDayNumber(strings.TrimSuffix(filepath.Base(path), ".jsonl"))
133 if err != nil {
134 return err
135 }
136
137 if day < place.day {
138 continue
139 }
140
141 offset := uint64(0)
142 if day == place.day {
143 offset = place.offset
144 }
145
146 err = journalScanFile(ctx, path, day, offset, n == len(paths)-1, add)
147 if err != nil {
148 return err
149 }
150 }
151
152 return nil
153 }
154
155 func (j *Journal) closeWriter() error {
156 if j.file == nil {
157 return nil
158 }
159
160 err := j.file.Close()
161 j.file = nil
162 j.day = 0
163 if err != nil {
164 return fmt.Errorf("close day file: %w", err)
165 }
166
167 return nil
168 }
169
170 // journalDayNumber returns the day YYYY-MM-DD as the number YYYYMMDD.
171 func journalDayNumber(day string) (uint32, error) {
172 number, err := strconv.ParseUint(strings.ReplaceAll(day, "-", ""), 10, 32)
173 if err != nil || len(day) != len("2006-01-02") {
174 return 0, fmt.Errorf("invalid day %q", day)
175 }
176
177 return uint32(number), nil
178 }
179
180 // journalDayName returns the day number YYYYMMDD as YYYY-MM-DD.
181 func journalDayName(day uint32) string {
182 return fmt.Sprintf("%04d-%02d-%02d", day/10000, day/100%100, day%100)
183 }
184
185 // journalScanFile gives each line of the day file at path from offset on to
186 // add. In the last day file, it cuts a part of a line at the end.
187 func journalScanFile(ctx context.Context, path string, day uint32, offset uint64, last bool, add func([]byte, Position) error) (err error) {
188 file, err := os.Open(path)
189 if err != nil {
190 return fmt.Errorf("open day file: %w", err)
191 }
192
193 defer func() {
194 err = errors.Join(err, file.Close())
195 }()
196
197 _, err = file.Seek(int64(offset), io.SeekStart)
198 if err != nil {
199 return fmt.Errorf("seek day file: %w", err)
200 }
201
202 reader := bufio.NewReader(file)
203 for {
204 if ctx.Err() != nil {
205 return ctx.Err()
206 }
207
208 line, err := reader.ReadBytes('\n')
209 if errors.Is(err, io.EOF) && len(line) == 0 {
210 return nil
211 }
212
213 if errors.Is(err, io.EOF) && last {
214 return journalCutFile(path, offset)
215 }
216
217 if errors.Is(err, io.EOF) {
218 return fmt.Errorf("%s ends in an incomplete line at %d", filepath.Base(path), offset)
219 }
220
221 if err != nil {
222 return fmt.Errorf("read %s: %w", filepath.Base(path), err)
223 }
224
225 place := Position{
226 day: day,
227 offset: offset,
228 length: uint32(len(line)),
229 }
230
231 err = add(line[:len(line)-1], place)
232 if err != nil {
233 return fmt.Errorf("%s at %d: %w", filepath.Base(path), offset, err)
234 }
235
236 offset += uint64(len(line))
237 }
238 }
239
240 // journalCutFile cuts the file at path to size bytes.
241 func journalCutFile(path string, size uint64) error {
242 err := os.Truncate(path, int64(size))
243 if err != nil {
244 return fmt.Errorf("cut incomplete line: %w", err)
245 }
246
247 return nil
248 }