feat: Allow connector to persist initial UID validity of mailboxes

Author: jameshoulahan

connector/connector.go

@@ -49,6 +49,12 @@ type Connector interface {
 	// It is recommended that the returned channel is buffered with at least constants.ChannelBufferCount.
 	GetUpdates() <-chan imap.Update
 
+	// GetUIDValidity returns the default UID validity for this user.
+	GetUIDValidity() imap.UID
+
+	// SetUIDValidity sets the default UID validity for this user.
+	SetUIDValidity(uidValidity imap.UID) error
+
 	// Close the connector will no longer be used and all resources should be closed/released.
 	Close(ctx context.Context) error
 }

connector/dummy.go

@@ -246,6 +246,14 @@ func (conn *Dummy) MarkMessagesFlagged(ctx context.Context, messageIDs []imap.Me
 	return nil
 }
 
+func (conn *Dummy) GetUIDValidity() imap.UID {
+	return 1
+}
+
+func (conn *Dummy) SetUIDValidity(imap.UID) error {
+	return nil
+}
+
 func (conn *Dummy) Sync(ctx context.Context) error {
 	for _, mailbox := range conn.state.getLabels() {
 		update := imap.NewMailboxCreated(mailbox)

internal/backend/connector_updates.go

@@ -94,7 +94,18 @@ func (user *user) applyMailboxDeleted(ctx context.Context, update *imap.MailboxD
 	}
 
 	return user.db.Write(ctx, func(ctx context.Context, tx *ent.Tx) error {
-		return db.DeleteMailboxWithRemoteID(ctx, tx, update.MailboxID)
+		uidValidity, increased, err := db.DeleteMailboxWithRemoteID(ctx, tx, update.MailboxID)
+		if err != nil {
+			return err
+		}
+
+		if increased {
+			if err := user.connector.SetUIDValidity(uidValidity); err != nil {
+				return err
+			}
+		}
+
+		return nil
 	})
 }
 

internal/backend/state_connector_impl.go

@@ -139,6 +139,10 @@ func (sc *stateConnectorImpl) SetMessagesFlagged(ctx context.Context, messageIDs
 	return nil
 }
 
+func (sc *stateConnectorImpl) SetUIDValidity(uidValidity imap.UID) error {
+	return sc.connector.SetUIDValidity(uidValidity)
+}
+
 func (sc *stateConnectorImpl) getMetadataValue(key string) any {
 	v, ok := sc.metadata[key]
 	if !ok {

internal/backend/user.go

@@ -42,7 +42,14 @@ type user struct {
 	messageIDCounter uint64
 }
 
-func newUser(ctx context.Context, userID string, database *db.DB, conn connector.Connector, store store.Store, delimiter string) (*user, error) {
+func newUser(
+	ctx context.Context,
+	userID string,
+	database *db.DB,
+	conn connector.Connector,
+	store store.Store,
+	delimiter string,
+) (*user, error) {
 	if err := database.Init(ctx); err != nil {
 		return nil, err
 	}
@@ -55,6 +62,13 @@ func newUser(ctx context.Context, userID string, database *db.DB, conn connector
 		return nil, err
 	}
 
+	// Init the global UID validity.
+	if err := database.Write(ctx, func(ctx context.Context, tx *ent.Tx) error {
+		return db.InitGlobalUIDValidity(ctx, tx, conn.GetUIDValidity())
+	}); err != nil {
+		return nil, err
+	}
+
 	user := &user{
 		userID:           userID,
 		connector:        conn,
@@ -79,10 +93,8 @@ func newUser(ctx context.Context, userID string, database *db.DB, conn connector
 		ctx, cancel := context.WithCancel(context.Background())
 		defer cancel()
 
-		ctx = contexts.NewRemoteUpdateCtx(ctx)
-
 		labels := pprof.Labels("go", "Connector Updates", "UserID", user.userID)
-		pprof.Do(ctx, labels, func(_ context.Context) {
+		pprof.Do(contexts.NewRemoteUpdateCtx(ctx), labels, func(_ context.Context) {
 			updateCh := user.updateInjector.GetUpdates()
 			for {
 				select {

internal/db/mailbox.go

@@ -70,27 +70,43 @@ func RenameMailboxWithRemoteID(ctx context.Context, tx *ent.Tx, mboxID imap.Labe
 	return nil
 }
 
-func DeleteMailboxWithRemoteID(ctx context.Context, tx *ent.Tx, mboxID imap.LabelID) error {
-	// get uid validity
+// DeleteMailboxWithRemoteID deletes the mailbox with the given remote ID.
+// It returns the (potentially new) global UID validity, along with a bool indicating whether it has been increased.
+func DeleteMailboxWithRemoteID(ctx context.Context, tx *ent.Tx, mboxID imap.LabelID) (imap.UID, bool, error) {
 	mbox, err := tx.Mailbox.Query().
 		Where(mailbox.RemoteID(mboxID)).
 		Select(mailbox.FieldUIDValidity).
 		Only(ctx)
 	if err != nil {
-		return err
+		return 0, false, err
 	}
 
-	if err := updateGlobalUIDValidity(ctx, tx, mbox.UIDValidity); err != nil {
-		return err
+	curUIDValidity, err := getGlobalUIDValidity(ctx, tx)
+	if err != nil {
+		return 0, false, err
+	}
+
+	var newUIDValidity imap.UID
+
+	if mbox.UIDValidity == curUIDValidity {
+		newUIDValidity = curUIDValidity.Add(1)
+	} else {
+		newUIDValidity = curUIDValidity
+	}
+
+	if newUIDValidity > curUIDValidity {
+		if err := setGlobalUIDValidity(ctx, tx, newUIDValidity); err != nil {
+			return 0, false, err
+		}
 	}
 
 	if _, err := tx.Mailbox.Delete().
 		Where(mailbox.RemoteID(mboxID)).
 		Exec(ctx); err != nil {
-		return err
+		return 0, false, err
 	}
 
-	return nil
+	return newUIDValidity, newUIDValidity > curUIDValidity, nil
 }
 
 func UpdateRemoteMailboxID(ctx context.Context, tx *ent.Tx, internalID imap.InternalMailboxID, remoteID imap.LabelID) error {
@@ -303,42 +319,42 @@ func FilterMailboxContains(ctx context.Context, client *ent.Client, mboxID imap.
 	return result, nil
 }
 
-func getGlobalUIDValidity(ctx context.Context, tx *ent.Tx) (imap.UID, error) {
-	globalUIDValidity, err := tx.UIDValidity.Query().
-		Select(uidvalidity.FieldUIDValidity).
-		Only(ctx)
+func InitGlobalUIDValidity(ctx context.Context, tx *ent.Tx, uidValidity imap.UID) error {
+	curUIDValidity, err := tx.UIDValidity.Query().Select(uidvalidity.FieldUIDValidity).Only(ctx)
+	if err != nil {
+		if !ent.IsNotFound(err) {
+			return fmt.Errorf("failed to get current UID validity: %w", err)
+		}
 
-	if ent.IsNotFound(err) {
-		_, err := tx.UIDValidity.Create().
-			SetUIDValidity(mailbox.DefaultUIDValidity).
-			Save(ctx)
-		if err != nil {
-			return 0, err
+		if _, err := tx.UIDValidity.Create().SetUIDValidity(uidValidity).Save(ctx); err != nil {
+			return fmt.Errorf("failed to create UID validity entry: %w", err)
 		}
 
-		return mailbox.DefaultUIDValidity, nil
+		return nil
 	}
 
-	if err != nil {
-		return 0, err
+	if uidValidity > curUIDValidity.UIDValidity {
+		if err := setGlobalUIDValidity(ctx, tx, uidValidity); err != nil {
+			return fmt.Errorf("failed to set UID validity: %w", err)
+		}
 	}
 
-	return globalUIDValidity.UIDValidity, nil
+	return nil
 }
 
-func updateGlobalUIDValidity(ctx context.Context, tx *ent.Tx, deletedUIDValidity imap.UID) error {
-	globalUIDValidity, err := getGlobalUIDValidity(ctx, tx)
+func getGlobalUIDValidity(ctx context.Context, tx *ent.Tx) (imap.UID, error) {
+	uidValidity, err := tx.UIDValidity.Query().Select(uidvalidity.FieldUIDValidity).Only(ctx)
 	if err != nil {
-		return err
+		return 0, err
 	}
 
-	if deletedUIDValidity < globalUIDValidity {
-		return nil
-	}
+	return uidValidity.UIDValidity, nil
+}
 
-	_, err = tx.UIDValidity.Update().
-		SetUIDValidity(deletedUIDValidity.Add(1)).
-		Save(ctx)
+func setGlobalUIDValidity(ctx context.Context, tx *ent.Tx, uidValidity imap.UID) error {
+	if _, err := tx.UIDValidity.Update().SetUIDValidity(uidValidity).Save(ctx); err != nil {
+		return err
+	}
 
-	return err
+	return nil
 }

internal/state/actions.go

@@ -61,10 +61,17 @@ func (state *State) actionDeleteMailbox(ctx context.Context, tx *ent.Tx, mboxID
 		return err
 	}
 
-	if err := db.DeleteMailboxWithRemoteID(ctx, tx, mboxID.RemoteID); err != nil {
+	uidValidity, increased, err := db.DeleteMailboxWithRemoteID(ctx, tx, mboxID.RemoteID)
+	if err != nil {
 		return err
 	}
 
+	if increased {
+		if err := state.user.GetRemote().SetUIDValidity(uidValidity); err != nil {
+			return err
+		}
+	}
+
 	return state.user.QueueOrApplyStateUpdate(ctx, tx, NewMailboxDeletedStateUpdate(mboxID.InternalID))
 }
 

internal/state/connector.go

@@ -69,4 +69,7 @@ type Connector interface {
 
 	// SetMessagesFlagged marks the message with the given ID as seen or unseen.
 	SetMessagesFlagged(ctx context.Context, messageIDs []imap.MessageID, flagged bool) error
+
+	// SetUIDValidity sets the UID Validity for the user.
+	SetUIDValidity(imap.UID) error
 }