Limit the file storage by an optional quota

1 files changed, 89 insertions, 0 deletions

diff --git a/xmpp/gateway/storage.go b/xmpp/gateway/storage.go
new file mode 100644
index 0000000..a3fad79
--- /dev/null
+++ b/xmpp/gateway/storage.go
@@ -0,0 +1,89 @@
+package gateway
+
+import (
+	"io/ioutil"
+	"os"
+	"sort"
+	"sync"
+
+	log "github.com/sirupsen/logrus"
+)
+
+// StorageQuota is a value from config parsed to bytes number
+var StorageQuota uint64
+// CachedStorageSize estimates the storage size between full rescans
+var CachedStorageSize uint64
+var StorageLock = sync.Mutex{}
+
+// MeasureStorageSize replaces the estimated storage size with relevant data from the filesystem
+func MeasureStorageSize(path string) {
+	dents, err := ioutil.ReadDir(path)
+	if err != nil {
+		return
+	}
+
+	var total uint64
+	for _, fi := range dents {
+		if !fi.IsDir() {
+			total += uint64(fi.Size())
+		}
+	}
+
+	if total != CachedStorageSize {
+		if CachedStorageSize > 0 {
+			log.Warnf("Correcting cached storage size: was %v, actually %v", CachedStorageSize, total)
+		}
+		CachedStorageSize = total
+	}
+}
+
+// CleanOldFiles purges the oldest files in a directory that exceed the limit
+func CleanOldFiles(path string, limit uint64) {
+	dents, err := ioutil.ReadDir(path)
+	if err != nil {
+		return
+	}
+
+	var total uint64
+	for _, fi := range dents {
+		if !fi.IsDir() {
+			total += uint64(fi.Size())
+		}
+	}
+
+	// sort by time
+	sort.Slice(dents, func(i int, j int) bool {
+		return dents[i].ModTime().Before(dents[j].ModTime())
+	})
+
+	// purge
+	if total > limit {
+		toPurge := total - limit
+
+		var purgedAmount uint64
+		var purgedCount uint64
+
+		for _, fi := range dents {
+			if !fi.IsDir() {
+				err = os.Remove(path + string(os.PathSeparator) + fi.Name())
+				if err != nil {
+					log.Errorf("Couldn't remove %v: %v", fi.Name(), err)
+					continue
+				}
+
+				purgedAmount += uint64(fi.Size())
+				purgedCount += 1
+				if purgedAmount >= toPurge {
+					break
+				}
+			}
+		}
+
+		log.Infof("Cleaned %v bytes of %v old files", purgedAmount, purgedCount)
+		if CachedStorageSize > purgedAmount {
+			CachedStorageSize -= purgedAmount
+		} else {
+			CachedStorageSize = 0
+		}
+	}
+}