Refine sync() to properly update sizes

message/cache.go

@@ -43,10 +43,10 @@ type queries struct {
 	selectAttachmentsExpired         string
 	selectAttachmentsSizeBySender    string
 	selectAttachmentsSizeByUserID    string
+	selectAttachmentsWithSizes       string
 	selectStats                      string
 	updateStats                      string
 	updateMessageTime                string
-	selectAttachmentIDs              string
 }
 
 // Cache stores published messages
@@ -363,16 +363,6 @@ func (c *Cache) ExpireMessages(topics ...string) error {
 	})
 }
 
-// AttachmentIDs returns message IDs with active (non-expired, non-deleted) attachments
-func (c *Cache) AttachmentIDs() ([]string, error) {
-	rows, err := c.db.ReadOnly().Query(c.queries.selectAttachmentIDs, time.Now().Unix())
-	if err != nil {
-		return nil, err
-	}
-	defer rows.Close()
-	return readStrings(rows)
-}
-
 // AttachmentsExpired returns message IDs with expired attachments that have not been deleted
 func (c *Cache) AttachmentsExpired() ([]string, error) {
 	rows, err := c.db.Query(c.queries.selectAttachmentsExpired, time.Now().Unix())
@@ -415,6 +405,30 @@ func (c *Cache) AttachmentBytesUsedByUser(userID string) (int64, error) {
 	return c.readAttachmentBytesUsed(rows)
 }
 
+// AttachmentsWithSizes returns a map of message ID to attachment size for all active
+// (non-expired, non-deleted) attachments. This is used to hydrate the attachment store's
+// size tracking on startup and during periodic sync.
+func (c *Cache) AttachmentsWithSizes() (map[string]int64, error) {
+	rows, err := c.db.ReadOnly().Query(c.queries.selectAttachmentsWithSizes, time.Now().Unix())
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	attachments := make(map[string]int64)
+	for rows.Next() {
+		var id string
+		var size int64
+		if err := rows.Scan(&id, &size); err != nil {
+			return nil, err
+		}
+		attachments[id] = size
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+	return attachments, nil
+}
+
 func (c *Cache) readAttachmentBytesUsed(rows *sql.Rows) (int64, error) {
 	defer rows.Close()
 	var size int64

message/cache_postgres.go

@@ -70,12 +70,11 @@ const (
 	postgresSelectAttachmentsExpiredQuery      = `SELECT mid FROM message WHERE attachment_expires > 0 AND attachment_expires <= $1 AND attachment_deleted = FALSE`
 	postgresSelectAttachmentsSizeBySenderQuery = `SELECT COALESCE(SUM(attachment_size), 0) FROM message WHERE user_id = '' AND sender = $1 AND attachment_expires >= $2`
 	postgresSelectAttachmentsSizeByUserIDQuery = `SELECT COALESCE(SUM(attachment_size), 0) FROM message WHERE user_id = $1 AND attachment_expires >= $2`
+	postgresSelectAttachmentsWithSizesQuery    = `SELECT mid, attachment_size FROM message WHERE attachment_expires > $1 AND attachment_deleted = FALSE`
 
 	postgresSelectStatsQuery       = `SELECT value FROM message_stats WHERE key = 'messages'`
 	postgresUpdateStatsQuery       = `UPDATE message_stats SET value = $1 WHERE key = 'messages'`
 	postgresUpdateMessageTimeQuery = `UPDATE message SET time = $1 WHERE mid = $2`
-
-	postgresSelectAttachmentIDsQuery = `SELECT mid FROM message WHERE attachment_expires > $1 AND attachment_deleted = FALSE`
 )
 
 var postgresQueries = queries{
@@ -99,10 +98,10 @@ var postgresQueries = queries{
 	selectAttachmentsExpired:         postgresSelectAttachmentsExpiredQuery,
 	selectAttachmentsSizeBySender:    postgresSelectAttachmentsSizeBySenderQuery,
 	selectAttachmentsSizeByUserID:    postgresSelectAttachmentsSizeByUserIDQuery,
+	selectAttachmentsWithSizes:       postgresSelectAttachmentsWithSizesQuery,
 	selectStats:                      postgresSelectStatsQuery,
 	updateStats:                      postgresUpdateStatsQuery,
 	updateMessageTime:                postgresUpdateMessageTimeQuery,
-	selectAttachmentIDs:              postgresSelectAttachmentIDsQuery,
 }
 
 // NewPostgresStore creates a new PostgreSQL-backed message cache store using an existing database connection pool.

message/cache_sqlite.go

@@ -73,12 +73,11 @@ const (
 	sqliteSelectAttachmentsExpiredQuery      = `SELECT mid FROM messages WHERE attachment_expires > 0 AND attachment_expires <= ? AND attachment_deleted = 0`
 	sqliteSelectAttachmentsSizeBySenderQuery = `SELECT IFNULL(SUM(attachment_size), 0) FROM messages WHERE user = '' AND sender = ? AND attachment_expires >= ?`
 	sqliteSelectAttachmentsSizeByUserIDQuery = `SELECT IFNULL(SUM(attachment_size), 0) FROM messages WHERE user = ? AND attachment_expires >= ?`
+	sqliteSelectAttachmentsWithSizesQuery    = `SELECT mid, attachment_size FROM messages WHERE attachment_expires > ? AND attachment_deleted = 0`
 
 	sqliteSelectStatsQuery       = `SELECT value FROM stats WHERE key = 'messages'`
 	sqliteUpdateStatsQuery       = `UPDATE stats SET value = ? WHERE key = 'messages'`
 	sqliteUpdateMessageTimeQuery = `UPDATE messages SET time = ? WHERE mid = ?`
-
-	sqliteSelectAttachmentIDsQuery = `SELECT mid FROM messages WHERE attachment_expires > ? AND attachment_deleted = 0`
 )
 
 var sqliteQueries = queries{
@@ -102,10 +101,10 @@ var sqliteQueries = queries{
 	selectAttachmentsExpired:         sqliteSelectAttachmentsExpiredQuery,
 	selectAttachmentsSizeBySender:    sqliteSelectAttachmentsSizeBySenderQuery,
 	selectAttachmentsSizeByUserID:    sqliteSelectAttachmentsSizeByUserIDQuery,
+	selectAttachmentsWithSizes:       sqliteSelectAttachmentsWithSizesQuery,
 	selectStats:                      sqliteSelectStatsQuery,
 	updateStats:                      sqliteUpdateStatsQuery,
 	updateMessageTime:                sqliteUpdateMessageTimeQuery,
-	selectAttachmentIDs:              sqliteSelectAttachmentIDsQuery,
 }
 
 // NewSQLiteStore creates a SQLite file-backed cache