1package store

3import (

4 "context"

5 "encoding/json"

6 "errors"

7 "fmt"

8 "io"

9 "log/slog"

10 "runtime"

11 "slices"

12 "sort"

13 "time"

15 "github.com/mjl-/bstore"

17 "github.com/mjl-/mox/message"

18 "github.com/mjl-/mox/mlog"

19 "github.com/mjl-/mox/moxio"

20)

22// Assign a new/incoming message to a thread. Message does not yet have an ID. If

23// this isn't a response, ThreadID should remain 0 (unless this is a message with

24// existing message-id) and the caller must set ThreadID to ID.

25// If the account is still busy upgrading messages with threadids in the background, parents

26// may have a threadid 0. That results in this message getting threadid 0, which

27// will handled by the background upgrade process assigning a threadid when it gets

28// to this message.

29func assignThread(log mlog.Log, tx *bstore.Tx, m *Message, part *message.Part) error {

30 if m.MessageID != "" {

31 // Match against existing different message with same Message-ID.

32 q := bstore.QueryTx[Message](tx)

33 q.FilterNonzero(Message{MessageID: m.MessageID})

34 q.FilterEqual("Expunged", false)

35 q.FilterNotEqual("ID", m.ID)

36 q.FilterNotEqual("ThreadID", int64(0))

37 q.SortAsc("ID")

38 q.Limit(1)

39 em, err := q.Get()

40 if err != nil && err != bstore.ErrAbsent {

41 return fmt.Errorf("looking up existing message with message-id: %v", err)

42 } else if err == nil {

43 assignParent(m, em, true)

44 return nil

45 }

46 }

48 h, err := part.Header()

49 if err != nil {

50 log.Errorx("assigning threads: parsing references/in-reply-to headers, not matching by message-id", err, slog.Int64("msgid", m.ID))

51 }

52 messageIDs, err := message.ReferencedIDs(h.Values("References"), h.Values("In-Reply-To"))

53 if err != nil {

54 log.Errorx("assigning threads: parsing references/in-reply-to headers, not matching by message-id", err, slog.Int64("msgid", m.ID))

55 }

56 for i := len(messageIDs) - 1; i >= 0; i-- {

57 messageID := messageIDs[i]

58 if messageID == m.MessageID {

59 continue

60 }

61 tm, _, err := lookupThreadMessage(tx, m.ID, messageID, m.SubjectBase, m.DSN)

62 if err != nil {

63 return fmt.Errorf("looking up thread message for new message: %v", err)

64 } else if tm != nil {

65 assignParent(m, *tm, true)

66 return nil

67 }

68 m.ThreadMissingLink = true

69 }

70 if len(messageIDs) > 0 {

71 return nil

72 }

74 var isResp bool

75 if part != nil && part.Envelope != nil {

76 m.SubjectBase, isResp = message.ThreadSubject(part.Envelope.Subject, false)

77 }

78 if !isResp || m.SubjectBase == "" {

79 return nil

80 }

81 m.ThreadMissingLink = true

82 tm, err := lookupThreadMessageSubject(tx, *m, m.SubjectBase)

83 if err != nil {

84 return fmt.Errorf("looking up thread message by subject: %v", err)

85 } else if tm != nil {

86 assignParent(m, *tm, true)

87 }

88 return nil

89}

91// assignParent assigns threading fields to m that make it a child of parent message pm.

92// updateSeen indicates if m.Seen should be cleared if pm is thread-muted.

93func assignParent(m *Message, pm Message, updateSeen bool) {

94 if pm.ThreadID == 0 {

95 panic(fmt.Sprintf("assigning message id %d/d%q to parent message id %d/%q which has threadid 0", m.ID, m.MessageID, pm.ID, pm.MessageID))

96 }

97 if m.ID == pm.ID {

98 panic(fmt.Sprintf("trying to make message id %d/%q its own parent", m.ID, m.MessageID))

99 }

100 m.ThreadID = pm.ThreadID

101 // Make sure we don't add cycles.

102 if !slices.Contains(pm.ThreadParentIDs, m.ID) {

103 m.ThreadParentIDs = append([]int64{pm.ID}, pm.ThreadParentIDs...)

104 } else if pm.ID != m.ID {

105 m.ThreadParentIDs = []int64{pm.ID}

106 } else {

107 m.ThreadParentIDs = nil

108 }

109 if m.MessageID != "" && m.MessageID == pm.MessageID {

110 m.ThreadMissingLink = true

111 }

112 m.ThreadMuted = pm.ThreadMuted

113 m.ThreadCollapsed = pm.ThreadCollapsed

114 if updateSeen && m.ThreadMuted {

115 m.Seen = true

116 }

117}

118

119// ResetThreading resets the MessageID and SubjectBase fields for all messages in

120// the account. If clearIDs is true, all Thread* fields are also cleared. Changes

121// are made in transactions of batchSize changes. The total number of updated

122// messages is returned.

123//

124// ModSeq is not changed. Callers should bump the uid validity of the mailboxes

125// to propagate the changes to IMAP clients.

126func (a *Account) ResetThreading(ctx context.Context, log mlog.Log, batchSize int, clearIDs bool) (int, error) {

127 // todo: should this send Change events for ThreadMuted and ThreadCollapsed? worth it?

128

129 var lastID int64

130 total := 0

131 for {

132 n := 0

133

134 prepareMessages := func(in, out chan moxio.Work[Message, Message]) {

135 for {

136 w, ok := <-in

137 if !ok {

138 return

139 }

140

141 m := w.In

142

143 // We have the Message-ID and Subject headers in ParsedBuf. We use a partial part

144 // struct so we don't generate so much garbage for the garbage collector to sift

145 // through.

146 var part struct {

147 Envelope *message.Envelope

148 }

149 if err := json.Unmarshal(m.ParsedBuf, &part); err != nil {

150 log.Errorx("unmarshal json parsedbuf for setting message-id, skipping", err, slog.Int64("msgid", m.ID))

151 } else {

152 m.MessageID = ""

153 if part.Envelope != nil && part.Envelope.MessageID != "" {

154 s, _, err := message.MessageIDCanonical(part.Envelope.MessageID)

155 if err != nil {

156 log.Debugx("parsing message-id, skipping", err, slog.Int64("msgid", m.ID), slog.String("messageid", part.Envelope.MessageID))

157 }

158 m.MessageID = s

159 }

160 if part.Envelope != nil {

161 m.SubjectBase, _ = message.ThreadSubject(part.Envelope.Subject, false)

162 }

163 }

164 w.Out = m

165

166 out <- w

167 }

168 }

169

170 err := a.DB.Write(ctx, func(tx *bstore.Tx) error {

171 processMessage := func(in, m Message) error {

172 if clearIDs {

173 m.ThreadID = 0

174 m.ThreadParentIDs = nil

175 m.ThreadMissingLink = false

176 }

177 return tx.Update(&m)

178 }

179

180 // JSON parsing is relatively heavy, we benefit from multiple goroutines.

181 procs := runtime.GOMAXPROCS(0)

182 wq := moxio.NewWorkQueue[Message, Message](procs, 2*procs, prepareMessages, processMessage)

183

184 q := bstore.QueryTx[Message](tx)

185 q.FilterEqual("Expunged", false)

186 q.FilterGreater("ID", lastID)

187 q.SortAsc("ID")

188 err := q.ForEach(func(m Message) error {

189 // We process in batches so we don't block other operations for a long time.

190 if n >= batchSize {

191 return bstore.StopForEach

192 }

193 // Update starting point for next batch.

194 lastID = m.ID

195

196 n++

197 return wq.Add(m)

198 })

199 if err == nil {

200 err = wq.Finish()

201 }

202 wq.Stop()

203 return err

204 })

205 if err != nil {

206 return total, fmt.Errorf("upgrading account to threads storage, step 1/2: %w", err)

207 }

208 total += n

209 if n == 0 {

210 break

211 }

212 }

213 return total, nil

214}

215

216// AssignThreads assigns thread-related fields to messages with ID >=

217// startMessageID. Changes are committed each batchSize changes if txOpt is nil

218// (i.e. during automatic account upgrade, we don't want to block database access

219// for a long time). If txOpt is not nil, all changes are made in that

220// transaction.

221//

222// When resetting thread assignments, the caller must first clear the existing

223// thread fields.

224//

225// Messages are processed in order of ID, so when added to the account, not

226// necessarily by received/date. Most threaded messages can immediately be matched

227// to their parent message. If not, we keep track of the missing message-id and

228// resolve as soon as we encounter it. At the end, we resolve all remaining

229// messages, they start with a cycle.

230//

231// Does not set Seen flag for muted threads.

232//

233// Progress is written to progressWriter, every 100k messages.

234func (a *Account) AssignThreads(ctx context.Context, log mlog.Log, txOpt *bstore.Tx, startMessageID int64, batchSize int, xprogressWriter io.Writer) error {

235 // We use a more basic version of the thread-matching algorithm describe in:

236 // ../rfc/5256:443

237 // The algorithm assumes you'll select messages, then group into threads. We normally do

238 // thread-calculation when messages are delivered. Here, we assign threads as soon

239 // as we can, but will queue messages that reference known ancestors and resolve as

240 // soon as we process them. We can handle large number of messages, but not very

241 // quickly because we make lots of database queries.

242

243 // xprogressWriter can call panic on write errors, when assigning threads through a

244 // ctl command.

245

246 type childMsg struct {

247 ID int64 // This message will be fetched and updated with the threading fields once the parent is resolved.

248 MessageID string // Of child message. Once child is resolved, its own children can be resolved too.

249 ThreadMissingLink bool

250 }

251 // Messages that have a References/In-Reply-To that we want to set as parent, but

252 // where the parent doesn't have a ThreadID yet are added to pending. The key is

253 // the normalized MessageID of the parent, and the value is a list of messages that

254 // can get resolved once the parent gets its ThreadID. The kids will get the same

255 // ThreadIDs, and they themselves may be parents to kids, and so on.

256 // For duplicate messages (messages with identical Message-ID), the second

257 // Message-ID to be added to pending is added under its own message-id, so it gets

258 // its original as parent.

259 pending := map[string][]childMsg{}

260

261 // Current tx. If not equal to txOpt, we clean it up before we leave.

262 var tx *bstore.Tx

263 defer func() {

264 if tx != nil && tx != txOpt {

265 err := tx.Rollback()

266 log.Check(err, "rolling back transaction")

267 }

268 }()

269

270 // Set thread-related fields for a single message. Caller must save the message,

271 // only if not an error and not added to the pending list.

272 assign := func(m *Message, references, inReplyTo []string, subject string) (pend bool, rerr error) {

273 if m.MessageID != "" {

274 // Attempt to match against existing different message with same Message-ID that

275 // already has a threadid.

276 // If there are multiple messages for a message-id a future call to assign may use

277 // its threadid, or it may end up in pending and we resolve it when we need to.

278 q := bstore.QueryTx[Message](tx)

279 q.FilterNonzero(Message{MessageID: m.MessageID})

280 q.FilterEqual("Expunged", false)

281 q.FilterLess("ID", m.ID)

282 q.SortAsc("ID")

283 q.Limit(1)

284 em, err := q.Get()

285 if err != nil && err != bstore.ErrAbsent {

286 return false, fmt.Errorf("looking up existing message with message-id: %v", err)

287 } else if err == nil {

288 if em.ThreadID == 0 {

289 pending[em.MessageID] = append(pending[em.MessageID], childMsg{m.ID, m.MessageID, true})

290 return true, nil

291 } else {

292 assignParent(m, em, false)

293 return false, nil

294 }

295 }

296 }

297

298 refids, err := message.ReferencedIDs(references, inReplyTo)

299 if err != nil {

300 log.Errorx("assigning threads: parsing references/in-reply-to headers, not matching by message-id", err, slog.Int64("msgid", m.ID))

301 }

302

303 for i := len(refids) - 1; i >= 0; i-- {

304 messageID := refids[i]

305 if messageID == m.MessageID {

306 continue

307 }

308 tm, exists, err := lookupThreadMessage(tx, m.ID, messageID, m.SubjectBase, m.DSN)

309 if err != nil {

310 return false, fmt.Errorf("lookup up thread by message-id %s for message id %d: %w", messageID, m.ID, err)

311 } else if tm != nil {

312 assignParent(m, *tm, false)

313 return false, nil

314 } else if exists {

315 pending[messageID] = append(pending[messageID], childMsg{m.ID, m.MessageID, i < len(refids)-1})

316 return true, nil

317 }

318 }

319

320 var subjectBase string

321 var isResp bool

322 if subject != "" {

323 subjectBase, isResp = message.ThreadSubject(subject, false)

324 }

325 if len(refids) > 0 || !isResp || subjectBase == "" {

326 m.ThreadID = m.ID

327 m.ThreadMissingLink = len(refids) > 0

328 return false, nil

329 }

330

331 // No references to use. If this is a reply/forward (based on subject), we'll match

332 // against base subject, at most 4 weeks back so we don't match against ancient

333 // messages and 1 day ahead so we can match against delayed deliveries.

334 tm, err := lookupThreadMessageSubject(tx, *m, subjectBase)

335 if err != nil {

336 return false, fmt.Errorf("looking up recent messages by base subject %q: %w", subjectBase, err)

337 } else if tm != nil {

338 m.ThreadID = tm.ThreadID

339 m.ThreadParentIDs = []int64{tm.ThreadID} // Always under root message with subject-match.

340 m.ThreadMissingLink = true

341 m.ThreadMuted = tm.ThreadMuted

342 m.ThreadCollapsed = tm.ThreadCollapsed

343 } else {

344 m.ThreadID = m.ID

345 }

346 return false, nil

347 }

348

349 npendingResolved := 0

350

351 // Resolve pending messages that wait on m.MessageID to be resolved, recursively.

352 var resolvePending func(tm Message, cyclic bool) error

353 resolvePending = func(tm Message, cyclic bool) error {

354 if tm.MessageID == "" {

355 return nil

356 }

357 l := pending[tm.MessageID]

358 delete(pending, tm.MessageID)

359 for _, mi := range l {

360 m := Message{ID: mi.ID}

361 if err := tx.Get(&m); err != nil {

362 return fmt.Errorf("get message %d for resolving pending thread for message-id %s, %d: %w", mi.ID, tm.MessageID, tm.ID, err)

363 } else if m.Expunged {

364 return fmt.Errorf("message %d marked as expunged", mi.ID)

365 }

366 if m.ThreadID != 0 {

367 // ThreadID already set because this is a cyclic message. If we would assign a

368 // parent again, we would create a cycle.

369 if m.MessageID != tm.MessageID && !cyclic {

370 panic(fmt.Sprintf("threadid already set (%d) while handling non-cyclic message id %d/%q and with different message-id %q as parent message id %d", m.ThreadID, m.ID, m.MessageID, tm.MessageID, tm.ID))

371 }

372 continue

373 }

374 assignParent(&m, tm, false)

375 m.ThreadMissingLink = mi.ThreadMissingLink

376 if err := tx.Update(&m); err != nil {

377 return fmt.Errorf("update message %d for resolving pending thread for message-id %s, %d: %w", mi.ID, tm.MessageID, tm.ID, err)

378 }

379 if err := resolvePending(m, cyclic); err != nil {

380 return err

381 }

382 npendingResolved++

383 }

384 return nil

385 }

386

387 // Output of the worker goroutines.

388 type threadPrep struct {

389 references []string

390 inReplyTo []string

391 subject string

392 }

393

394 // Single allocation.

395 threadingFields := [][]byte{

396 []byte("references"),

397 []byte("in-reply-to"),

398 []byte("subject"),

399 }

400

401 // Worker goroutine function. We start with a reasonably large buffer for reading

402 // the header into. And we have scratch space to copy the needed headers into. That

403 // means we normally won't allocate any more buffers.

404 prepareMessages := func(in, out chan moxio.Work[Message, threadPrep]) {

405 headerbuf := make([]byte, 8*1024)

406 scratch := make([]byte, 4*1024)

407 for {

408 w, ok := <-in

409 if !ok {

410 return

411 }

412

413 m := w.In

414 var partialPart struct {

415 HeaderOffset int64

416 BodyOffset int64

417 }

418 if err := json.Unmarshal(m.ParsedBuf, &partialPart); err != nil {

419 w.Err = fmt.Errorf("unmarshal part: %v", err)

420 } else {

421 size := partialPart.BodyOffset - partialPart.HeaderOffset

422 if int(size) > len(headerbuf) {

423 headerbuf = make([]byte, size)

424 }

425 if size > 0 {

426 buf := headerbuf[:int(size)]

427 err := func() error {

428 mr := a.MessageReader(m)

429 defer func() {

430 err := mr.Close()

431 log.Check(err, "closing message reader")

432 }()

433

434 // ReadAt returns whole buffer or error. Single read should be fast.

435 n, err := mr.ReadAt(buf, partialPart.HeaderOffset)

436 if err != nil || n != len(buf) {

437 return fmt.Errorf("read header: %v", err)

438 }

439 return nil

440 }()

441 if err != nil {

442 w.Err = err

443 } else if h, err := message.ParseHeaderFields(buf, scratch, threadingFields); err != nil {

444 w.Err = err

445 } else {

446 w.Out.references = h["References"]

447 w.Out.inReplyTo = h["In-Reply-To"]

448 l := h["Subject"]

449 if len(l) > 0 {

450 w.Out.subject = l[0]

451 }

452 }

453 }

454 }

455

456 out <- w

457 }

458 }

459

460 // Assign threads to messages, possibly in batches.

461 nassigned := 0

462 for {

463 n := 0

464 tx = txOpt

465 if tx == nil {

466 var err error

467 tx, err = a.DB.Begin(ctx, true)

468 if err != nil {

469 return fmt.Errorf("begin transaction: %w", err)

470 }

471 }

472

473 processMessage := func(m Message, prep threadPrep) error {

474 pend, err := assign(&m, prep.references, prep.inReplyTo, prep.subject)

475 if err != nil {

476 return fmt.Errorf("for msgid %d: %w", m.ID, err)

477 } else if pend {

478 return nil

479 }

480 if m.ThreadID == 0 {

481 panic(fmt.Sprintf("no threadid after assign of message id %d/%q", m.ID, m.MessageID))

482 }

483 // Fields have been set, store in database and resolve messages waiting for this MessageID.

484 if slices.Contains(m.ThreadParentIDs, m.ID) {

485 panic(fmt.Sprintf("message id %d/%q contains itself in parent ids %v", m.ID, m.MessageID, m.ThreadParentIDs))

486 }

487 if err := tx.Update(&m); err != nil {

488 return err

489 }

490 if err := resolvePending(m, false); err != nil {

491 return fmt.Errorf("resolving pending message-id: %v", err)

492 }

493 return nil

494 }

495

496 // Use multiple worker goroutines to parse headers from on-disk messages.

497 procs := runtime.GOMAXPROCS(0)

498 wq := moxio.NewWorkQueue[Message, threadPrep](2*procs, 4*procs, prepareMessages, processMessage)

499

500 // We assign threads in order by ID, so messages delivered in between our

501 // transaction will get assigned threads too: they'll have the highest id's.

502 q := bstore.QueryTx[Message](tx)

503 q.FilterGreaterEqual("ID", startMessageID)

504 q.FilterEqual("Expunged", false)

505 q.SortAsc("ID")

506 err := q.ForEach(func(m Message) error {

507 // Batch number of changes, so we give other users of account a change to run.

508 if txOpt == nil && n >= batchSize {

509 return bstore.StopForEach

510 }

511 // Starting point for next batch.

512 startMessageID = m.ID + 1

513 // Don't process again. Can happen when earlier upgrade was aborted.

514 if m.ThreadID != 0 {

515 return nil

516 }

517

518 n++

519 return wq.Add(m)

520 })

521 if err == nil {

522 err = wq.Finish()

523 }

524 wq.Stop()

525

526 if err == nil && txOpt == nil {

527 err = tx.Commit()

528 tx = nil

529 }

530 if err != nil {

531 return fmt.Errorf("assigning threads: %w", err)

532 }

533 if n == 0 {

534 break

535 }

536 nassigned += n

537 if nassigned%100000 == 0 {

538 log.Debug("assigning threads, progress", slog.Int("count", nassigned), slog.Int("unresolved", len(pending)))

539 if _, err := fmt.Fprintf(xprogressWriter, "assigning threads, progress: %d messages\n", nassigned); err != nil {

540 return fmt.Errorf("writing progress: %v", err)

541 }

542 }

543 }

544 if _, err := fmt.Fprintf(xprogressWriter, "assigning threads, done: %d messages\n", nassigned); err != nil {

545 return fmt.Errorf("writing progress: %v", err)

546 }

547

548 log.Debug("assigning threads, mostly done, finishing with resolving of cyclic messages", slog.Int("count", nassigned), slog.Int("unresolved", len(pending)))

549

550 if _, err := fmt.Fprintf(xprogressWriter, "assigning threads, resolving %d cyclic pending message-ids\n", len(pending)); err != nil {

551 return fmt.Errorf("writing progress: %v", err)

552 }

553

554 // Remaining messages in pending have cycles and possibly tails. The cycle is at

555 // the head of the thread. Once we resolve that, the rest of the thread can be

556 // resolved too. Ignoring self-references (duplicate messages), there can only be

557 // one cycle, and it is at the head. So we look for cycles, ignoring

558 // self-references, and resolve a message as soon as we see the cycle.

559

560 parent := map[string]string{} // Child Message-ID pointing to the parent Message-ID, excluding self-references.