]> git.giacomolli.uk - onechat/blob - journal.go
Talk to the agent through a Telegram bot
[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 := dayNumber(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 func (j *journal) closeWriter() error {
80 if j.file == nil {
81 return nil
82 }
83
84 err := j.file.Close()
85 j.file = nil
86 j.day = 0
87 if err != nil {
88 return fmt.Errorf("close day file: %w", err)
89 }
90
91 return nil
92 }
93
94 // Read returns the line at place without its newline.
95 func (j *journal) Read(place position) ([]byte, error) {
96 reader, open := j.readers[place.day]
97 if !open {
98 var err error
99 reader, err = os.Open(filepath.Join(j.dir, dayName(place.day)+".jsonl"))
100 if err != nil {
101 return nil, fmt.Errorf("open day file: %w", err)
102 }
103
104 if j.readers == nil {
105 j.readers = make(map[uint32]*os.File)
106 }
107
108 j.readers[place.day] = reader
109 }
110
111 line := make([]byte, place.length)
112 _, err := reader.ReadAt(line, int64(place.offset))
113 if err != nil {
114 return nil, fmt.Errorf("read line of %s: %w", dayName(place.day), err)
115 }
116
117 if len(line) == 0 || line[len(line)-1] != '\n' {
118 return nil, fmt.Errorf("line of %s at %d does not end its record", dayName(place.day), place.offset)
119 }
120
121 return line[:len(line)-1], nil
122 }
123
124 // Close closes the open day files.
125 func (j *journal) Close() error {
126 err := j.closeWriter()
127 for day, reader := range j.readers {
128 err = errors.Join(err, reader.Close())
129 delete(j.readers, day)
130 }
131
132 return err
133 }
134
135 // Scan gives each line of the day files from place on to add, with its
136 // position, oldest day first. The zero position starts at the first line. A
137 // crash during a write can leave a part of a line at the end of the last day
138 // file. Its write never returned, so Scan cuts it off.
139 func (j *journal) Scan(ctx context.Context, place position, add func([]byte, position) error) error {
140 paths, err := filepath.Glob(filepath.Join(j.dir, "*.jsonl"))
141 if err != nil {
142 return fmt.Errorf("list day files: %w", err)
143 }
144
145 slices.Sort(paths)
146 for n, path := range paths {
147 day, err := dayNumber(strings.TrimSuffix(filepath.Base(path), ".jsonl"))
148 if err != nil {
149 return err
150 }
151
152 if day < place.day {
153 continue
154 }
155
156 offset := uint64(0)
157 if day == place.day {
158 offset = place.offset
159 }
160
161 err = scanFile(ctx, path, day, offset, n == len(paths)-1, add)
162 if err != nil {
163 return err
164 }
165 }
166
167 return nil
168 }
169
170 // scanFile gives each line of the day file at path from offset on to add. In
171 // the last day file, it cuts a part of a line at the end.
172 func scanFile(ctx context.Context, path string, day uint32, offset uint64, last bool, add func([]byte, position) error) (err error) {
173 file, err := os.Open(path)
174 if err != nil {
175 return fmt.Errorf("open day file: %w", err)
176 }
177
178 defer func() {
179 err = errors.Join(err, file.Close())
180 }()
181
182 _, err = file.Seek(int64(offset), io.SeekStart)
183 if err != nil {
184 return fmt.Errorf("seek day file: %w", err)
185 }
186
187 reader := bufio.NewReader(file)
188 for {
189 if ctx.Err() != nil {
190 return ctx.Err()
191 }
192
193 line, err := reader.ReadBytes('\n')
194 if errors.Is(err, io.EOF) && len(line) == 0 {
195 return nil
196 }
197
198 if errors.Is(err, io.EOF) && last {
199 return cutFile(path, offset)
200 }
201
202 if errors.Is(err, io.EOF) {
203 return fmt.Errorf("%s ends in an incomplete line at %d", filepath.Base(path), offset)
204 }
205
206 if err != nil {
207 return fmt.Errorf("read %s: %w", filepath.Base(path), err)
208 }
209
210 place := position{
211 day: day,
212 offset: offset,
213 length: uint32(len(line)),
214 }
215
216 err = add(line[:len(line)-1], place)
217 if err != nil {
218 return fmt.Errorf("%s at %d: %w", filepath.Base(path), offset, err)
219 }
220
221 offset += uint64(len(line))
222 }
223 }
224
225 // cutFile cuts the file at path to size bytes.
226 func cutFile(path string, size uint64) error {
227 err := os.Truncate(path, int64(size))
228 if err != nil {
229 return fmt.Errorf("cut incomplete line: %w", err)
230 }
231
232 return nil
233 }
234
235 // dayNumber returns the day YYYY-MM-DD as the number YYYYMMDD.
236 func dayNumber(day string) (uint32, error) {
237 number, err := strconv.ParseUint(strings.ReplaceAll(day, "-", ""), 10, 32)
238 if err != nil || len(day) != len("2006-01-02") {
239 return 0, fmt.Errorf("invalid day %q", day)
240 }
241
242 return uint32(number), nil
243 }
244
245 // dayName returns the day number YYYYMMDD as YYYY-MM-DD.
246 func dayName(day uint32) string {
247 return fmt.Sprintf("%04d-%02d-%02d", day/10000, day/100%100, day%100)
248 }