forked from mirrors/ntfy
Expire attachments properly
This commit is contained in:
parent
c45a28e6af
commit
e7c19a2bad
5 changed files with 101 additions and 14 deletions
|
@ -21,4 +21,5 @@ type cache interface {
|
|||
Prune(olderThan time.Time) error
|
||||
MarkPublished(m *message) error
|
||||
AttachmentsSize(owner string) (int64, error)
|
||||
AttachmentsExpired() ([]string, error)
|
||||
}
|
||||
|
|
|
@ -139,6 +139,20 @@ func (c *memCache) AttachmentsSize(owner string) (int64, error) {
|
|||
return size, nil
|
||||
}
|
||||
|
||||
func (c *memCache) AttachmentsExpired() ([]string, error) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
ids := make([]string, 0)
|
||||
for topic := range c.messages {
|
||||
for _, m := range c.messages[topic] {
|
||||
if m.Attachment != nil && m.Attachment.Expires > 0 && m.Attachment.Expires < time.Now().Unix() {
|
||||
ids = append(ids, m.ID)
|
||||
}
|
||||
}
|
||||
}
|
||||
return ids, nil
|
||||
}
|
||||
|
||||
func (c *memCache) pruneTopic(topic string, olderThan time.Time) {
|
||||
messages := make([]*message, 0)
|
||||
for _, m := range c.messages[topic] {
|
||||
|
|
|
@ -60,7 +60,8 @@ const (
|
|||
selectMessagesCountQuery = `SELECT COUNT(*) FROM messages`
|
||||
selectMessageCountForTopicQuery = `SELECT COUNT(*) FROM messages WHERE topic = ?`
|
||||
selectTopicsQuery = `SELECT topic FROM messages GROUP BY topic`
|
||||
selectAttachmentsSizeQuery = `SELECT IFNULL(SUM(attachment_size), 0) FROM messages WHERE attachment_owner = ?`
|
||||
selectAttachmentsSizeQuery = `SELECT IFNULL(SUM(attachment_size), 0) FROM messages WHERE attachment_owner = ? AND attachment_expires >= ?`
|
||||
selectAttachmentsExpiredQuery = `SELECT id FROM messages WHERE attachment_expires > 0 AND attachment_expires < ?`
|
||||
)
|
||||
|
||||
// Schema management queries
|
||||
|
@ -234,7 +235,7 @@ func (c *sqliteCache) Prune(olderThan time.Time) error {
|
|||
}
|
||||
|
||||
func (c *sqliteCache) AttachmentsSize(owner string) (int64, error) {
|
||||
rows, err := c.db.Query(selectAttachmentsSizeQuery, owner)
|
||||
rows, err := c.db.Query(selectAttachmentsSizeQuery, owner, time.Now().Unix())
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
@ -251,6 +252,26 @@ func (c *sqliteCache) AttachmentsSize(owner string) (int64, error) {
|
|||
return size, nil
|
||||
}
|
||||
|
||||
func (c *sqliteCache) AttachmentsExpired() ([]string, error) {
|
||||
rows, err := c.db.Query(selectAttachmentsExpiredQuery, time.Now().Unix())
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
ids := make([]string, 0)
|
||||
for rows.Next() {
|
||||
var id string
|
||||
if err := rows.Scan(&id); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ids = append(ids, id)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return ids, nil
|
||||
}
|
||||
|
||||
func readMessages(rows *sql.Rows) ([]*message, error) {
|
||||
defer rows.Close()
|
||||
messages := make([]*message, 0)
|
||||
|
|
|
@ -28,18 +28,10 @@ func newFileCache(dir string, totalSizeLimit int64, fileSizeLimit int64) (*fileC
|
|||
if err := os.MkdirAll(dir, 0700); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
entries, err := os.ReadDir(dir)
|
||||
size, err := dirSize(dir)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var size int64
|
||||
for _, e := range entries {
|
||||
info, err := e.Info()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
size += info.Size()
|
||||
}
|
||||
return &fileCache{
|
||||
dir: dir,
|
||||
totalSizeCurrent: size,
|
||||
|
@ -58,8 +50,8 @@ func (c *fileCache) Write(id string, in io.Reader, limiters ...*util.Limiter) (i
|
|||
return 0, err
|
||||
}
|
||||
defer f.Close()
|
||||
log.Printf("remaining total: %d", c.remainingTotalSize())
|
||||
limiters = append(limiters, util.NewLimiter(c.remainingTotalSize()), util.NewLimiter(c.fileSizeLimit))
|
||||
log.Printf("remaining total: %d", c.Remaining())
|
||||
limiters = append(limiters, util.NewLimiter(c.Remaining()), util.NewLimiter(c.fileSizeLimit))
|
||||
limitWriter := util.NewLimitWriter(f, limiters...)
|
||||
size, err := io.Copy(limitWriter, in)
|
||||
if err != nil {
|
||||
|
@ -77,7 +69,40 @@ func (c *fileCache) Write(id string, in io.Reader, limiters ...*util.Limiter) (i
|
|||
|
||||
}
|
||||
|
||||
func (c *fileCache) remainingTotalSize() int64 {
|
||||
func (c *fileCache) Remove(ids []string) error {
|
||||
var firstErr error
|
||||
for _, id := range ids {
|
||||
if err := c.removeFile(id); err != nil {
|
||||
if firstErr == nil {
|
||||
firstErr = err // Continue despite error; we want to delete as many as we can
|
||||
}
|
||||
}
|
||||
}
|
||||
size, err := dirSize(c.dir)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
c.mu.Lock()
|
||||
c.totalSizeCurrent = size
|
||||
c.mu.Unlock()
|
||||
return firstErr
|
||||
}
|
||||
|
||||
func (c *fileCache) removeFile(id string) error {
|
||||
if !fileIDRegex.MatchString(id) {
|
||||
return errInvalidFileID
|
||||
}
|
||||
file := filepath.Join(c.dir, id)
|
||||
return os.Remove(file)
|
||||
}
|
||||
|
||||
func (c *fileCache) Size() int64 {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
return c.totalSizeCurrent
|
||||
}
|
||||
|
||||
func (c *fileCache) Remaining() int64 {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
remaining := c.totalSizeLimit - c.totalSizeCurrent
|
||||
|
@ -86,3 +111,19 @@ func (c *fileCache) remainingTotalSize() int64 {
|
|||
}
|
||||
return remaining
|
||||
}
|
||||
|
||||
func dirSize(dir string) (int64, error) {
|
||||
entries, err := os.ReadDir(dir)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
var size int64
|
||||
for _, e := range entries {
|
||||
info, err := e.Info()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
size += info.Size()
|
||||
}
|
||||
return size, nil
|
||||
}
|
||||
|
|
|
@ -832,6 +832,16 @@ func (s *Server) updateStatsAndPrune() {
|
|||
}
|
||||
}
|
||||
|
||||
// Delete expired attachments
|
||||
ids, err := s.cache.AttachmentsExpired()
|
||||
if err == nil {
|
||||
if err := s.fileCache.Remove(ids); err != nil {
|
||||
log.Printf("error while deleting attachments: %s", err.Error())
|
||||
}
|
||||
} else {
|
||||
log.Printf("error retrieving expired attachments: %s", err.Error())
|
||||
}
|
||||
|
||||
// Prune message cache
|
||||
olderThan := time.Now().Add(-1 * s.config.CacheDuration)
|
||||
if err := s.cache.Prune(olderThan); err != nil {
|
||||
|
|
Loading…
Reference in a new issue