Skip to content

Commit c08f510

Browse files
LBeernaertProtonLBeernaertProton
committed
fix(GODT-1570): Fix too many SQL Variable in sync benchmark
Batch incoming connector updates for new message to 1000 messages per transaction to avoid running into too many SQL variables. We batched at outside of the transaction level so as not to completely block the other threads forever until we are done processing a high number of updates. The code has also been updated to use a more optimized database query.
1 parent 8c9b3f1 commit c08f510

3 files changed

Lines changed: 156 additions & 55 deletions

File tree

‎internal/backend/connector_updates.go‎

Lines changed: 64 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -159,76 +159,89 @@ func (user *user) applyMessagesCreated(ctx context.Context, update *imap.Message
159159
return err
160160
}
161161

162-
var reqs []*db.CreateMessageReq
162+
remoteMessageRequestsMap := make(map[imap.MessageID]*db.CreateMessageReq, len(updates))
163163

164-
remoteToLocalMessageID := make(map[imap.MessageID]imap.InternalMessageID)
164+
const maxUpdateChunk = 1000
165165

166-
return db.WriteAndStore(ctx, user.db, user.store, func(ctx context.Context, tx *ent.Tx, storeTx store.Transaction) error {
167-
for _, update := range updates {
168-
internalID := uuid.NewString()
166+
chunkedUpdates := xslices.Chunk(updates, maxUpdateChunk)
169167

170-
literal, err := rfc822.SetHeaderValue(update.Literal, ids.InternalIDKey, internalID)
171-
if err != nil {
172-
return fmt.Errorf("failed to set internal ID: %w", err)
173-
}
168+
for _, chunk := range chunkedUpdates {
169+
messagesForMailboxes := make(map[imap.InternalMailboxID][]*db.CreateMessageReq)
174170

175-
if err := storeTx.Set(imap.InternalMessageID(internalID), literal); err != nil {
176-
return fmt.Errorf("failed to store message literal: %w", err)
177-
}
171+
if err := db.WriteAndStore(ctx, user.db, user.store, func(ctx context.Context, tx *ent.Tx, storeTx store.Transaction) error {
172+
for _, update := range chunk {
178173

179-
reqs = append(reqs, &db.CreateMessageReq{
180-
Message: update.Message,
181-
Literal: literal,
182-
Body: update.Body,
183-
Structure: update.Structure,
184-
Envelope: update.Envelope,
185-
InternalID: imap.InternalMessageID(internalID),
186-
})
174+
var request *db.CreateMessageReq
187175

188-
remoteToLocalMessageID[update.Message.ID] = imap.InternalMessageID(internalID)
189-
}
176+
// Do not reprocess the message if we have previously processed it.
177+
if req, ok := remoteMessageRequestsMap[update.Message.ID]; ok {
178+
request = req
179+
} else {
190180

191-
if _, err := db.CreateMessages(ctx, tx, reqs...); err != nil {
192-
return fmt.Errorf("failed to create message: %w", err)
193-
}
181+
internalID := uuid.NewString()
194182

195-
messageIDs := make(map[imap.LabelID][]ids.MessageIDPair)
183+
literal, err := rfc822.SetHeaderValue(update.Literal, ids.InternalIDKey, internalID)
184+
if err != nil {
185+
return fmt.Errorf("failed to set internal ID: %w", err)
186+
}
196187

197-
for _, update := range updates {
198-
for _, mailboxID := range update.MailboxIDs {
199-
localID := remoteToLocalMessageID[update.Message.ID]
200-
idPair := ids.MessageIDPair{
201-
InternalID: localID,
202-
RemoteID: update.Message.ID,
203-
}
188+
if err := storeTx.Set(imap.InternalMessageID(internalID), literal); err != nil {
189+
return fmt.Errorf("failed to store message literal: %w", err)
190+
}
191+
192+
request = &db.CreateMessageReq{
193+
Message: update.Message,
194+
Literal: literal,
195+
Body: update.Body,
196+
Structure: update.Structure,
197+
Envelope: update.Envelope,
198+
InternalID: imap.InternalMessageID(internalID),
199+
}
204200

205-
if !slices.Contains(messageIDs[mailboxID], idPair) {
206-
messageIDs[mailboxID] = append(messageIDs[mailboxID], idPair)
201+
remoteMessageRequestsMap[update.Message.ID] = req
207202
}
208-
}
209-
}
210203

211-
for mailboxID, messageIDs := range messageIDs {
212-
internalMailboxID, err := db.GetMailboxIDWithRemoteID(ctx, tx.Client(), mailboxID)
213-
if err != nil {
214-
return err
215-
}
204+
for _, mbox := range update.MailboxIDs {
205+
internalMailboxID, err := db.GetMailboxIDWithRemoteID(ctx, tx.Client(), mbox)
206+
if err != nil {
207+
return err
208+
}
216209

217-
internalIDs := xslices.Map(messageIDs, func(id ids.MessageIDPair) imap.InternalMessageID {
218-
return id.InternalID
219-
})
210+
existingRequests := messagesForMailboxes[internalMailboxID]
220211

221-
messageUIDs, err := db.AddMessagesToMailbox(ctx, tx, internalIDs, internalMailboxID)
222-
if err != nil {
223-
return err
212+
var alreadyPresent bool
213+
214+
for _, e := range existingRequests {
215+
if e.Message.ID == request.Message.ID {
216+
alreadyPresent = true
217+
break
218+
}
219+
}
220+
221+
if alreadyPresent {
222+
continue
223+
}
224+
225+
messagesForMailboxes[internalMailboxID] = append(messagesForMailboxes[internalMailboxID], request)
226+
}
224227
}
225228

226-
user.queueStateUpdate(state.NewExistsStateUpdate(internalMailboxID, messageIDs, messageUIDs, nil))
229+
for mbox, request := range messagesForMailboxes {
230+
result, err := db.CreateAndAddMessagesToMailbox(ctx, tx, mbox, request)
231+
if err != nil {
232+
return err
233+
}
234+
235+
user.queueStateUpdate(state.NewExistsStateUpdate(mbox, result, nil))
236+
}
227237

238+
return nil
239+
}); err != nil {
240+
return err
228241
}
242+
}
229243

230-
return nil
231-
})
244+
return nil
232245
}
233246

234247
// applyMessageLabelsUpdated applies a MessageLabelsUpdated update.

‎internal/db/message.go‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import (
99
"github.com/ProtonMail/gluon/internal/db/ent/message"
1010
"github.com/ProtonMail/gluon/internal/db/ent/messageflag"
1111
"github.com/ProtonMail/gluon/internal/db/ent/uid"
12+
"github.com/ProtonMail/gluon/internal/ids"
1213
"github.com/bradenaw/juniper/xslices"
1314
)
1415

@@ -131,6 +132,7 @@ func CreateAndAddMessageToMailbox(ctx context.Context, tx *ent.Tx, mboxID imap.I
131132
SetMessage(message).
132133
SetUID(mbox.UIDNext).
133134
Save(ctx)
135+
134136
if err != nil {
135137
return 0, imap.FlagSet{}, err
136138
}
@@ -142,6 +144,93 @@ func CreateAndAddMessageToMailbox(ctx context.Context, tx *ent.Tx, mboxID imap.I
142144
return uid.UID, NewFlagSet(uid, flags), err
143145
}
144146

147+
type CreateAndAddMessagesResult struct {
148+
UID imap.UID
149+
Flags imap.FlagSet
150+
MessageID ids.MessageIDPair
151+
}
152+
153+
func CreateAndAddMessagesToMailbox(ctx context.Context, tx *ent.Tx, mboxID imap.InternalMailboxID, requests []*CreateMessageReq) ([]CreateAndAddMessagesResult, error) {
154+
mbox, err := tx.Mailbox.Query().Where(mailbox.ID(mboxID)).Select(mailbox.FieldID, mailbox.FieldUIDNext).Only(ctx)
155+
if err != nil {
156+
return nil, err
157+
}
158+
159+
msgBuilders := make([]*ent.MessageCreate, 0, len(requests))
160+
flags := make([][]*ent.MessageFlag, 0, len(requests))
161+
162+
for _, request := range requests {
163+
builders := xslices.Map(request.Message.Flags.ToSlice(), func(flag string) *ent.MessageFlagCreate {
164+
return tx.MessageFlag.Create().SetValue(flag)
165+
})
166+
167+
entFlags, err := tx.MessageFlag.CreateBulk(builders...).Save(ctx)
168+
if err != nil {
169+
return nil, err
170+
}
171+
172+
flags = append(flags, entFlags)
173+
174+
msgCreate := tx.Message.Create().
175+
SetID(request.InternalID).
176+
SetDate(request.Message.Date).
177+
SetBody(request.Body).
178+
SetBodyStructure(request.Structure).
179+
SetEnvelope(request.Envelope).
180+
SetSize(len(request.Literal)).
181+
AddFlags(entFlags...)
182+
183+
if len(request.Message.ID) != 0 {
184+
msgCreate = msgCreate.SetRemoteID(request.Message.ID)
185+
}
186+
187+
msgBuilders = append(msgBuilders, msgCreate)
188+
}
189+
190+
messages, err := tx.Message.CreateBulk(msgBuilders...).Save(ctx)
191+
if err != nil {
192+
return nil, err
193+
}
194+
195+
uidBuilders := make([]*ent.UIDCreate, 0, len(requests))
196+
197+
for i, message := range messages {
198+
uidBuilders = append(uidBuilders, tx.UID.Create().
199+
SetMailboxID(mbox.ID).
200+
SetMessageID(message.ID).
201+
SetUID(mbox.UIDNext.Add(uint32(i))),
202+
)
203+
}
204+
205+
uids, err := tx.UID.CreateBulk(uidBuilders...).Save(ctx)
206+
if err != nil {
207+
return nil, err
208+
}
209+
210+
if err := BumpMailboxUIDNext(ctx, tx, mbox, len(requests)); err != nil {
211+
return nil, err
212+
}
213+
214+
result := make([]CreateAndAddMessagesResult, 0, len(requests))
215+
216+
for i := 0; i < len(requests); i++ {
217+
if uids[i].UID != mbox.UIDNext.Add(uint32(i)) {
218+
panic("Invalid UID ")
219+
}
220+
221+
result = append(result, CreateAndAddMessagesResult{
222+
MessageID: ids.MessageIDPair{
223+
InternalID: messages[i].ID,
224+
RemoteID: messages[i].RemoteID,
225+
},
226+
UID: uids[i].UID,
227+
Flags: NewFlagSet(uids[i], flags[i]),
228+
})
229+
}
230+
231+
return result, err
232+
}
233+
145234
func BumpMailboxUIDsForMessage(ctx context.Context, tx *ent.Tx, messageIDs []imap.InternalMessageID, mboxID imap.InternalMailboxID) (map[imap.InternalMessageID]*ent.UID, error) {
146235
messageUIDs := make(map[imap.InternalMessageID]imap.UID)
147236

‎internal/state/responders.go‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -149,7 +149,7 @@ type ExistsStateUpdate struct {
149149
targetStateSet bool
150150
}
151151

152-
func NewExistsStateUpdate(mailboxID imap.InternalMailboxID, messageIDs []ids.MessageIDPair, uids map[imap.InternalMessageID]*ent.UID, s *State) Update {
152+
func NewExistsStateUpdate(mailboxID imap.InternalMailboxID, messages []db.CreateAndAddMessagesResult, s *State) Update {
153153
var stateID StateID
154154

155155
var targetStateSet bool
@@ -159,9 +159,8 @@ func NewExistsStateUpdate(mailboxID imap.InternalMailboxID, messageIDs []ids.Mes
159159
targetStateSet = true
160160
}
161161

162-
responders := xslices.Map(messageIDs, func(messageID ids.MessageIDPair) *exists {
163-
uid := uids[messageID.InternalID]
164-
exists := newExists(ids.NewMessageIDPair(uid.Edges.Message), uid.UID, db.NewFlagSet(uid, uid.Edges.Message.Edges.Flags))
162+
responders := xslices.Map(messages, func(result db.CreateAndAddMessagesResult) *exists {
163+
exists := newExists(result.MessageID, result.UID, result.Flags)
165164

166165
return exists
167166
})

0 commit comments

Comments
 (0)