1package store

3import (

4 "context"

5 "encoding/json"

6 "errors"

7 "fmt"

8 "io"

9 "runtime"

10 "sort"

11 "time"

13 "golang.org/x/exp/slices"

14 "golang.org/x/exp/slog"

16 "github.com/mjl-/bstore"

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

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

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

21)

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

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

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

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

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

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

29// to this message.

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

31 if m.MessageID != "" {

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

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

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

35 q.FilterEqual("Expunged", false)

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

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

38 q.SortAsc("ID")

39 q.Limit(1)

40 em, err := q.Get()

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

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

43 } else if err == nil {

44 assignParent(m, em, true)

45 return nil

46 }

47 }

49 h, err := part.Header()

50 if err != nil {

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

52 }

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

54 if err != nil {

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

56 }

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

58 messageID := messageIDs[i]

59 if messageID == m.MessageID {

60 continue

61 }

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

63 if err != nil {

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

65 } else if tm != nil {

66 assignParent(m, *tm, true)

67 return nil

68 }

69 m.ThreadMissingLink = true

70 }

71 if len(messageIDs) > 0 {

72 return nil

73 }

75 var isResp bool

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

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

78 }

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

80 return nil

81 }

82 m.ThreadMissingLink = true

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

84 if err != nil {

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

86 } else if tm != nil {

87 assignParent(m, *tm, true)

88 }

89 return nil

90}

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

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

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

95 if pm.ThreadID == 0 {

96 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))

97 }

98 if m.ID == pm.ID {

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

100 }

101 m.ThreadID = pm.ThreadID

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

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

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

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

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

107 } else {

108 m.ThreadParentIDs = nil

109 }

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

111 m.ThreadMissingLink = true

112 }

113 m.ThreadMuted = pm.ThreadMuted

114 m.ThreadCollapsed = pm.ThreadCollapsed

115 if updateSeen && m.ThreadMuted {

116 m.Seen = true

117 }

118}

119

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

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

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

123// messages is returned.

124//

125// ModSeq is not changed. Calles should bump the uid validity of the mailboxes

126// to propagate the changes to IMAP clients.

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

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

129

130 var lastID int64

131 total := 0

132 for {

133 n := 0

134

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

136 for {

137 w, ok := <-in

138 if !ok {

139 return

140 }

141

142 m := w.In

143

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

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

146 // through.

147 var part struct {

148 Envelope *message.Envelope

149 }

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

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

152 } else {

153 m.MessageID = ""

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

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

156 if err != nil {

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

158 }

159 m.MessageID = s

160 }

161 if part.Envelope != nil {

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

163 }

164 }

165 w.Out = m

166

167 out <- w

168 }

169 }

170

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

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

173 if clearIDs {

174 m.ThreadID = 0

175 m.ThreadParentIDs = nil

176 m.ThreadMissingLink = false

177 }

178 return tx.Update(&m)

179 }

180

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

182 procs := runtime.GOMAXPROCS(0)

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

184

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

186 q.FilterEqual("Expunged", false)

187 q.FilterGreater("ID", lastID)

188 q.SortAsc("ID")

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

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

191 if n >= batchSize {

192 return bstore.StopForEach

193 }

194 // Update starting point for next batch.

195 lastID = m.ID

196

197 n++

198 return wq.Add(m)

199 })

200 if err == nil {

201 err = wq.Finish()

202 }

203 wq.Stop()

204 return err

205 })

206 if err != nil {

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

208 }

209 total += n

210 if n == 0 {

211 break

212 }

213 }

214 return total, nil

215}

216

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

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

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

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

221// transaction.

222//

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

224// thread fields.

225//

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

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

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

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

230// messages, they start with a cycle.

231//

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

233//

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

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

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

237 // ../rfc/5256:443

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

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

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

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

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

243

244 type childMsg struct {

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

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

247 ThreadMissingLink bool

248 }

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

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

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

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

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

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

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

256 // its original as parent.

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

258

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

260 var tx *bstore.Tx

261 defer func() {

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

263 err := tx.Rollback()

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

265 }

266 }()

267

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

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

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

271 if m.MessageID != "" {

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

273 // already has a threadid.

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

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

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

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

278 q.FilterEqual("Expunged", false)

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

280 q.SortAsc("ID")

281 q.Limit(1)

282 em, err := q.Get()

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

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

285 } else if err == nil {

286 if em.ThreadID == 0 {

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

288 return true, nil

289 } else {

290 assignParent(m, em, false)

291 return false, nil

292 }

293 }

294 }

295

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

297 if err != nil {

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

299 }

300

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

302 messageID := refids[i]

303 if messageID == m.MessageID {

304 continue

305 }

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

307 if err != nil {

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

309 } else if tm != nil {

310 assignParent(m, *tm, false)

311 return false, nil

312 } else if exists {

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

314 return true, nil

315 }

316 }

317

318 var subjectBase string

319 var isResp bool

320 if subject != "" {

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

322 }

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

324 m.ThreadID = m.ID

325 m.ThreadMissingLink = len(refids) > 0

326 return false, nil

327 }

328

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

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

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

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

333 if err != nil {

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

335 } else if tm != nil {

336 m.ThreadID = tm.ThreadID

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

338 m.ThreadMissingLink = true

339 m.ThreadMuted = tm.ThreadMuted

340 m.ThreadCollapsed = tm.ThreadCollapsed

341 } else {

342 m.ThreadID = m.ID

343 }

344 return false, nil

345 }

346

347 npendingResolved := 0

348

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

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

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

352 if tm.MessageID == "" {

353 return nil

354 }

355 l := pending[tm.MessageID]

356 delete(pending, tm.MessageID)

357 for _, mi := range l {

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

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

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

361 }

362 if m.ThreadID != 0 {

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

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

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

366 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))

367 }

368 continue

369 }

370 assignParent(&m, tm, false)

371 m.ThreadMissingLink = mi.ThreadMissingLink

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

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

374 }

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

376 return err

377 }

378 npendingResolved++

379 }

380 return nil

381 }

382

383 // Output of the worker goroutines.

384 type threadPrep struct {

385 references []string

386 inReplyTo []string

387 subject string

388 }

389

390 // Single allocation.

391 threadingFields := [][]byte{

392 []byte("references"),

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

394 []byte("subject"),

395 }

396

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

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

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

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

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

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

403 for {

404 w, ok := <-in

405 if !ok {

406 return

407 }

408

409 m := w.In

410 var partialPart struct {

411 HeaderOffset int64

412 BodyOffset int64

413 }

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

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

416 } else {

417 size := partialPart.BodyOffset - partialPart.HeaderOffset

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

419 headerbuf = make([]byte, size)

420 }

421 if size > 0 {

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

423 err := func() error {

424 mr := a.MessageReader(m)

425 defer mr.Close()

426

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

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

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

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

431 }

432 return nil

433 }()

434 if err != nil {

435 w.Err = err

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

437 w.Err = err

438 } else {

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

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

441 l := h["Subject"]

442 if len(l) > 0 {

443 w.Out.subject = l[0]

444 }

445 }

446 }

447 }

448

449 out <- w

450 }

451 }

452

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

454 nassigned := 0

455 for {

456 n := 0

457 tx = txOpt

458 if tx == nil {

459 var err error

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

461 if err != nil {

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

463 }

464 }

465

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

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

468 if err != nil {

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

470 } else if pend {

471 return nil

472 }

473 if m.ThreadID == 0 {

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

475 }

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

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

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

479 }

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

481 return err

482 }

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

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

485 }

486 return nil

487 }

488

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

490 procs := runtime.GOMAXPROCS(0)

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

492

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

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

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

496 q.FilterGreaterEqual("ID", startMessageID)

497 q.FilterEqual("Expunged", false)

498 q.SortAsc("ID")

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

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

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

502 return bstore.StopForEach

503 }

504 // Starting point for next batch.

505 startMessageID = m.ID + 1

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

507 if m.ThreadID != 0 {

508 return nil

509 }

510

511 n++

512 return wq.Add(m)

513 })

514 if err == nil {

515 err = wq.Finish()

516 }

517 wq.Stop()

518

519 if err == nil && txOpt == nil {

520 err = tx.Commit()

521 tx = nil

522 }

523 if err != nil {

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

525 }

526 if n == 0 {

527 break

528 }

529 nassigned += n

530 if nassigned%100000 == 0 {

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

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

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

534 }

535 }

536 }

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

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

539 }

540

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

542

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

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

545 }

546

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

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

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

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

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

552

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