Add body read timeout recovery to IMAP client

Introduce Client wrapper with socket deadline support
Add reconnection logic for body read timeouts
Implement test cases for underflow scenarios
Update orchestrator to handle reconnections
This commit is contained in:
2026-07-21 05:29:35 +07:00
parent c741cd19a0
commit 6bf3a6c4ca
6 changed files with 375 additions and 30 deletions
+78 -16
View File
@@ -8,6 +8,7 @@ import (
"io"
"log/slog"
"net"
"os"
"strings"
"time"
@@ -42,8 +43,31 @@ type CopyDeps struct {
// in-body signal the watchdog can't tell a slow-but-live transfer from a
// wedged connection and cancels a healthy copy. May be nil.
OnActivity func()
// ReconnectSrc dials and logs in a FRESH source client, returning it. It is
// called when a body read times out (server under-delivered a literal),
// which leaves the current src connection desynced and unusable. CopyFolder
// swaps to the returned client, re-EXAMINEs the folder, and resumes. The
// implementation is expected to also update any external reference to the
// live src client (e.g. so a cancel path closes the right connection). If
// nil, a body-read timeout aborts the folder instead of recovering.
ReconnectSrc func() (*Client, error)
}
// ErrBodyTimeout means a message body did not finish transferring within the
// idle deadline — the server stopped sending mid-literal. It is almost always a
// server announcing a BODY[] literal larger than the bytes it actually sends,
// which makes go-imap wait forever for bytes that never come. Distinct from a
// closed connection so the caller can skip just this one message and resume.
var ErrBodyTimeout = errors.New("message body read timed out")
// bodyIdleTimeout bounds how long a body read may go with NO bytes arriving
// before it is abandoned. It is generous enough for legitimately slow servers
// (even ~15 KB/s links keep bytes flowing far more often than this) yet well
// under the orchestrator's multi-minute stall watchdog, so an under-delivered
// literal is caught quickly and locally instead of stalling the whole account.
// A var (not const) so tests can shorten it.
var bodyIdleTimeout = 30 * time.Second
// CopyResult summarizes the outcome of one CopyFolder run.
type CopyResult struct {
Copied int
@@ -95,7 +119,7 @@ func metaBatches(total, batchSize uint32) []imap.SeqRange {
// held in memory only for the duration of a single FETCH->APPEND and is
// never written to disk. Messages already migrated (per deps.IsMigrated)
// are skipped without re-fetching their bodies.
func CopyFolder(ctx context.Context, src, dst *imapclient.Client, srcFolder, dstFolder string, deps CopyDeps) (CopyResult, error) {
func CopyFolder(ctx context.Context, src, dst *Client, srcFolder, dstFolder string, deps CopyDeps) (CopyResult, error) {
var res CopyResult
sel, err := src.Select(srcFolder, &imap.SelectOptions{ReadOnly: true}).Wait()
@@ -190,6 +214,25 @@ func CopyFolder(ctx context.Context, src, dst *imapclient.Client, srcFolder, dst
if err := streamOne(src, dst, dstFolder, q.uid, q.flags, q.internalDate, deps.OnActivity); err != nil {
res.Errors++
reportErr(msgRef(q.uid, q.subject), "copy message: "+err.Error())
// A body-read timeout means the server under-delivered this message's
// literal; the src connection is now desynced. Mark the message
// migrated so this and future runs skip it (it is un-fetchable via a
// conforming client), then reconnect src and resume the folder with
// the remaining queued messages.
if errors.Is(err, ErrBodyTimeout) && deps.ReconnectSrc != nil {
if merr := deps.MarkMigrated(dstFolder, q.key); merr != nil {
reportErr(msgRef(q.uid, q.subject), "mark skipped: "+merr.Error())
}
newSrc, rerr := deps.ReconnectSrc()
if rerr != nil {
return res, fmt.Errorf("reconnect src after body timeout in %q: %w", srcFolder, rerr)
}
src = newSrc
if _, serr := src.Select(srcFolder, &imap.SelectOptions{ReadOnly: true}).Wait(); serr != nil {
return res, fmt.Errorf("re-examine %q after reconnect: %w", srcFolder, serr)
}
continue
}
// A closed/broken connection won't recover: every remaining APPEND
// would fail identically. Abort the folder instead of logging
// thousands of the same error; a re-run resumes via dedup.
@@ -218,17 +261,26 @@ func CopyFolder(ctx context.Context, src, dst *imapclient.Client, srcFolder, dst
// evidence of which side and which message is the culprit.
const slowMessage = 20 * time.Second
// touchReader wraps a body Read stream and pings onActivity on every non-empty
// read, so a long FETCH keeps the stall watchdog fed byte-by-byte.
type touchReader struct {
r io.Reader
on func()
// deadlineReader arms an idle read deadline on the source socket before every
// read of a message body, so a server that stops sending mid-literal — having
// announced a larger BODY[] size than it actually delivers — trips the deadline
// instead of blocking go-imap forever waiting for bytes that never arrive. It
// also pings onActivity as bytes arrive, feeding the orchestrator's stall
// watchdog. The deadline is refreshed on each read, so it bounds IDLE time
// (no bytes) rather than total transfer time — a legitimately slow but steady
// download never trips it.
type deadlineReader struct {
c *Client
r io.Reader
idle time.Duration
on func()
}
func (t touchReader) Read(p []byte) (int, error) {
n, err := t.r.Read(p)
if n > 0 && t.on != nil {
t.on()
func (d deadlineReader) Read(p []byte) (int, error) {
_ = d.c.SetReadDeadline(time.Now().Add(d.idle))
n, err := d.r.Read(p)
if n > 0 && d.on != nil {
d.on()
}
return n, err
}
@@ -252,7 +304,11 @@ func (t touchWriter) Write(p []byte) (int, error) {
// spooling to disk. The body is buffered in RAM only for the duration of
// this single FETCH->APPEND round trip. onActivity (may be nil) fires as bytes
// move in either direction, feeding the orchestrator's stall watchdog.
func streamOne(src, dst *imapclient.Client, dstFolder string, uid imap.UID, flags []imap.Flag, internalDate time.Time, onActivity func()) error {
//
// The body read is guarded by an idle deadline on src: if the server goes
// silent mid-literal, streamOne returns ErrBodyTimeout rather than hanging, so
// CopyFolder can skip the message and reconnect.
func streamOne(src, dst *Client, dstFolder string, uid imap.UID, flags []imap.Flag, internalDate time.Time, onActivity func()) error {
bodySection := &imap.FetchItemBodySection{}
fetchStart := time.Now()
fetchCmd := src.Fetch(imap.UIDSetNum(uid), &imap.FetchOptions{
@@ -265,19 +321,25 @@ func streamOne(src, dst *imapclient.Client, dstFolder string, uid imap.UID, flag
return fmt.Errorf("no message for uid %v", uid)
}
var body []byte
var readErr error
for {
item := msg.Next()
if item == nil {
break
}
if d, ok := item.(imapclient.FetchItemDataBodySection); ok {
b, err := io.ReadAll(touchReader{r: d.Literal, on: onActivity})
if err != nil {
return err
}
body = b
body, readErr = io.ReadAll(deadlineReader{c: src, r: d.Literal, idle: bodyIdleTimeout, on: onActivity})
}
}
// Clear the deadline before any further I/O on src (fetchCmd.Close reads the
// command's completion off the same socket).
_ = src.SetReadDeadline(time.Time{})
if readErr != nil {
if errors.Is(readErr, os.ErrDeadlineExceeded) {
return fmt.Errorf("%w: uid %v (server sent fewer bytes than the announced literal)", ErrBodyTimeout, uid)
}
return readErr
}
if err := fetchCmd.Close(); err != nil {
return err
}