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