diff --git a/internal/user/sync.go b/internal/user/sync.go index 76f9b61e..215c3b82 100644 --- a/internal/user/sync.go +++ b/internal/user/sync.go @@ -503,41 +503,7 @@ func syncMessages( return } - var expectedMemUsage uint64 - var chunks [][]proton.FullMessage - - { - var lastIndex int - var index int - for _, v := range buildBatch.batch { - var dataSize uint64 - for _, a := range v.Attachments { - dataSize += uint64(a.Size) - } - - // 2x increase for attachment due to extra memory needed for decrypting and writing - // in memory buffer. - dataSize *= 2 - dataSize += uint64(len(v.Body)) - - nextMemSize := expectedMemUsage + dataSize - if nextMemSize >= syncMaxMessageBuildingMem { - chunks = append(chunks, buildBatch.batch[lastIndex:index]) - lastIndex = index - expectedMemUsage = dataSize - } else { - expectedMemUsage = nextMemSize - } - - index++ - } - - if index < len(buildBatch.batch) { - chunks = append(chunks, buildBatch.batch[index:]) - } else if index == len(buildBatch.batch) && len(chunks) == 0 { - chunks = [][]proton.FullMessage{buildBatch.batch} - } - } + chunks := chunkSyncBuilderBatch(buildBatch.batch, syncMaxMessageBuildingMem) for index, chunk := range chunks { logrus.Debugf("Build request: %v of %v count=%v", index, len(chunks), len(chunk)) @@ -852,3 +818,39 @@ func (a *attachmentDownloader) getAttachments(ctx context.Context, attachments [ func (a *attachmentDownloader) close() { a.cancel() } + +func chunkSyncBuilderBatch(batch []proton.FullMessage, maxMemory uint64) [][]proton.FullMessage { + var expectedMemUsage uint64 + var chunks [][]proton.FullMessage + var lastIndex int + var index int + + for _, v := range batch { + var dataSize uint64 + for _, a := range v.Attachments { + dataSize += uint64(a.Size) + } + + // 2x increase for attachment due to extra memory needed for decrypting and writing + // in memory buffer. + dataSize *= 2 + dataSize += uint64(len(v.Body)) + + nextMemSize := expectedMemUsage + dataSize + if nextMemSize >= maxMemory { + chunks = append(chunks, batch[lastIndex:index]) + lastIndex = index + expectedMemUsage = dataSize + } else { + expectedMemUsage = nextMemSize + } + + index++ + } + + if lastIndex < len(batch) { + chunks = append(chunks, batch[lastIndex:]) + } + + return chunks +} diff --git a/internal/user/sync_build_test.go b/internal/user/sync_build_test.go index f52f5ef9..3f0aff2e 100644 --- a/internal/user/sync_build_test.go +++ b/internal/user/sync_build_test.go @@ -24,6 +24,8 @@ import ( "github.com/ProtonMail/gluon/imap" "github.com/ProtonMail/gluon/rfc822" + "github.com/ProtonMail/go-proton-api" + "github.com/bradenaw/juniper/xslices" "github.com/stretchr/testify/require" ) @@ -47,3 +49,32 @@ func TestNewFailedMessageLiteral(t *testing.T) { require.Equal(t, `("text" "plain" () NIL NIL "base64" 114 2)`, parsed.Body) require.Equal(t, `("text" "plain" () NIL NIL "base64" 114 2 NIL NIL NIL NIL)`, parsed.Structure) } + +func TestSyncChunkSyncBuilderBatch(t *testing.T) { + // GODT-2424 - Some messages were not fully built due to a bug in the chunking if the total memory used by the + // message would be higher than the maximum we allowed. + const totalMessageCount = 100 + + msg := proton.FullMessage{ + Message: proton.Message{ + Attachments: []proton.Attachment{ + { + Size: int64(8 * Megabyte), + }, + }, + }, + AttData: nil, + } + + messages := xslices.Repeat(msg, totalMessageCount) + + chunks := chunkSyncBuilderBatch(messages, 16*Megabyte) + + var totalMessagesInChunks int + + for _, v := range chunks { + totalMessagesInChunks += len(v) + } + + require.Equal(t, totalMessagesInChunks, totalMessageCount) +}