headers: enable partial header fetching

Enable partial header fetching by creating config values for headers to
specifically include, or specifically exclude. The References field will
always be fetched, regardless of the include list. Envelope data is
always fetched, but is not shown with :toggle-headers, since it isn't in
the RFC822 struct unless explicitly included in the list.

Partial headers can break the cache on changes. Update the cache tag key
to include the state of the partially-fetched headers.

Partial header fetching can have a significant performance increase for
IMAP, and for all backends a resource improvement. Some data to support
this is below. Gathered by opening aerc, selecting a mailbox with
approximately 800 messages and scrolling to the end.

Received measured with nethogs, RAM from btop

		   Received  |   RAM
-------------------------------------
All Headers	|   9,656 kb | 103 MB
Minimum Headers	|     896 kb |  36 MB

Signed-off-by: Tim Culverhouse <tim@timculverhouse.com>
Acked-by: Moritz Poldrack <moritz@poldrack.dev>
Acked-by: Robin Jarry <robin@jarry.cc>

8 files changed, 80 insertions, 5 deletions

diff --git a/worker/imap/cache.go b/worker/imap/cache.go
index a02f2cae..cf372154 100644
--- a/worker/imap/cache.go
+++ b/worker/imap/cache.go
@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"os"
 	"path"
+	"strings"
 	"time"
 
 	"git.sr.ht/~rjarry/aerc/lib/parse"
@@ -42,6 +43,14 @@ var (
 // initCacheDb opens (or creates) the database for the cache. One database is
 // created per account
 func (w *IMAPWorker) initCacheDb(acct string) {
+	switch {
+	case len(w.config.headersExclude) > 0:
+		headerTag := strings.Join(w.config.headersExclude, "")
+		cacheTag = append(cacheTag, headerTag...)
+	case len(w.config.headers) > 0:
+		headerTag := strings.Join(w.config.headers, "")
+		cacheTag = append(cacheTag, headerTag...)
+	}
 	cd, err := cacheDir()
 	if err != nil {
 		w.cache = nil
diff --git a/worker/imap/configure.go b/worker/imap/configure.go
index a9689f68..1581794f 100644
--- a/worker/imap/configure.go
+++ b/worker/imap/configure.go
@@ -60,6 +60,8 @@ func (w *IMAPWorker) handleConfigure(msg *types.Configure) error {
 
 	w.config.user = u.User
 	w.config.folders = msg.Config.Folders
+	w.config.headers = msg.Config.Headers
+	w.config.headersExclude = msg.Config.HeadersExclude
 
 	w.config.idle_timeout = 10 * time.Second
 	w.config.idle_debounce = 10 * time.Millisecond
diff --git a/worker/imap/fetch.go b/worker/imap/fetch.go
index f209d7f0..8189e6aa 100644
--- a/worker/imap/fetch.go
+++ b/worker/imap/fetch.go
@@ -29,11 +29,19 @@ func (imapw *IMAPWorker) handleFetchMessageHeaders(
 		return
 	}
 	log.Tracef("Fetching message headers: %v", toFetch)
+	hdrBodyPart := imap.BodyPartName{
+		Specifier: imap.HeaderSpecifier,
+	}
+	switch {
+	case len(imapw.config.headersExclude) > 0:
+		hdrBodyPart.NotFields = true
+		hdrBodyPart.Fields = imapw.config.headersExclude
+	case len(imapw.config.headers) > 0:
+		hdrBodyPart.Fields = imapw.config.headers
+	}
 	section := &imap.BodySectionName{
-		BodyPartName: imap.BodyPartName{
-			Specifier: imap.HeaderSpecifier,
-		},
-		Peek: true,
+		BodyPartName: hdrBodyPart,
+		Peek:         true,
 	}
 
 	items := []imap.FetchItem{
diff --git a/worker/imap/worker.go b/worker/imap/worker.go
index f9a722e6..f46f39d6 100644
--- a/worker/imap/worker.go
+++ b/worker/imap/worker.go
@@ -43,6 +43,8 @@ type imapConfig struct {
 	insecure          bool
 	addr              string
 	user              *url.Userinfo
+	headers           []string
+	headersExclude    []string
 	folders           []string
 	oauthBearer       lib.OAuthBearer
 	xoauth2           lib.Xoauth2
diff --git a/worker/lib/parse.go b/worker/lib/parse.go
index ac86292f..2baf13f7 100644
--- a/worker/lib/parse.go
+++ b/worker/lib/parse.go
@@ -310,6 +310,30 @@ func MessageInfo(raw RawMessage) (*models.MessageInfo, error) {
 	}, nil
 }
 
+// LimitHeaders returns a new Header with the specified headers included or
+// excluded
+func LimitHeaders(hdr *mail.Header, fields []string, exclude bool) {
+	fieldMap := make(map[string]struct{}, len(fields))
+	for _, f := range fields {
+		fieldMap[strings.ToLower(f)] = struct{}{}
+	}
+	curFields := hdr.Fields()
+	for curFields.Next() {
+		key := strings.ToLower(curFields.Key())
+		_, ok := fieldMap[key]
+		switch {
+		case exclude && ok:
+			curFields.Del()
+		case exclude && !ok:
+			// No-op: exclude but we didn't find it
+		case !exclude && ok:
+			// No-op: include and we found it
+		case !exclude && !ok:
+			curFields.Del()
+		}
+	}
+}
+
 // MessageHeaders populates a models.MessageInfo struct for the message.
 // based on the reader returned by NewReader. Minimal information is included.
 // There is no body structure or RFC822Headers set
diff --git a/worker/maildir/worker.go b/worker/maildir/worker.go
index b222aab8..c45992eb 100644
--- a/worker/maildir/worker.go
+++ b/worker/maildir/worker.go
@@ -47,6 +47,8 @@ type Worker struct {
 	currentSortCriteria []*types.SortCriterion
 	maildirpp           bool // whether to use Maildir++ directory layout
 	capabilities        *models.Capabilities
+	headers             []string
+	headersExclude      []string
 }
 
 // NewWorker creates a new maildir worker with the provided worker.
@@ -349,6 +351,8 @@ func (w *Worker) handleConfigure(msg *types.Configure) error {
 	if err != nil {
 		return err
 	}
+	w.headers = msg.Config.Headers
+	w.headersExclude = msg.Config.HeadersExclude
 	log.Debugf("configured base maildir: %s", dir)
 	return nil
 }
@@ -620,6 +624,12 @@ func (w *Worker) handleFetchMessageHeaders(
 			w.worker.PostMessageInfoError(msg, uid, err)
 			continue
 		}
+		switch {
+		case len(w.headersExclude) > 0:
+			lib.LimitHeaders(info.RFC822Headers, w.headersExclude, true)
+		case len(w.headers) > 0:
+			lib.LimitHeaders(info.RFC822Headers, w.headers, false)
+		}
 		w.worker.PostMessage(&types.MessageInfo{
 			Message: types.RespondTo(msg),
 			Info:    info,
diff --git a/worker/mbox/worker.go b/worker/mbox/worker.go
index 47713048..7a746cf0 100644
--- a/worker/mbox/worker.go
+++ b/worker/mbox/worker.go
@@ -30,7 +30,9 @@ type mboxWorker struct {
 	folder *container
 	worker *types.Worker
 
-	capabilities *models.Capabilities
+	capabilities   *models.Capabilities
+	headers        []string
+	headersExclude []string
 }
 
 func NewWorker(worker *types.Worker) (types.Backend, error) {
@@ -68,6 +70,8 @@ func (w *mboxWorker) handleMessage(msg types.WorkerMessage) error {
 		} else {
 			dir = filepath.Join(u.Host, u.Path)
 		}
+		w.headers = msg.Config.Headers
+		w.headersExclude = msg.Config.HeadersExclude
 		w.data, err = createMailboxContainer(dir)
 		if err != nil || w.data == nil {
 			w.data = &mailboxContainer{
@@ -161,6 +165,12 @@ func (w *mboxWorker) handleMessage(msg types.WorkerMessage) error {
 				w.worker.PostMessageInfoError(msg, uid, err)
 				break
 			} else {
+				switch {
+				case len(w.headersExclude) > 0:
+					lib.LimitHeaders(msgInfo.RFC822Headers, w.headersExclude, true)
+				case len(w.headers) > 0:
+					lib.LimitHeaders(msgInfo.RFC822Headers, w.headers, false)
+				}
 				w.worker.PostMessage(&types.MessageInfo{
 					Message: types.RespondTo(msg),
 					Info:    msgInfo,
diff --git a/worker/notmuch/worker.go b/worker/notmuch/worker.go
index d5c012ad..a7bab688 100644
--- a/worker/notmuch/worker.go
+++ b/worker/notmuch/worker.go
@@ -49,6 +49,8 @@ type worker struct {
 	watcher             types.FSWatcher
 	watcherDebounce     *time.Timer
 	capabilities        *models.Capabilities
+	headers             []string
+	headersExclude      []string
 }
 
 // NewWorker creates a new notmuch worker with the provided worker.
@@ -223,6 +225,8 @@ func (w *worker) handleConfigure(msg *types.Configure) error {
 		}
 		w.store = store
 	}
+	w.headers = msg.Config.Headers
+	w.headersExclude = msg.Config.HeadersExclude
 
 	return nil
 }
@@ -676,6 +680,12 @@ func (w *worker) emitMessageInfo(m *Message,
 	if err != nil {
 		return fmt.Errorf("could not get MessageInfo: %w", err)
 	}
+	switch {
+	case len(w.headersExclude) > 0:
+		lib.LimitHeaders(info.RFC822Headers, w.headersExclude, true)
+	case len(w.headers) > 0:
+		lib.LimitHeaders(info.RFC822Headers, w.headers, false)
+	}
 	w.w.PostMessage(&types.MessageInfo{
 		Message: types.RespondTo(parent),
 		Info:    info,