Limit the file storage by an optional quota

This commit is contained in:
Bohdan Horbeshko 2022-03-31 21:42:12 -04:00
parent 5c238db1da
commit 17afd3f8c7
8 changed files with 246 additions and 1 deletions

View file

@ -4,7 +4,7 @@ all:
go build -o telegabber go build -o telegabber
test: test:
go test -v ./config ./ ./telegram ./xmpp/gateway ./persistence ./telegram/formatter go test -v ./config ./ ./telegram ./xmpp ./xmpp/gateway ./persistence ./telegram/formatter
lint: lint:
$(GOPATH)/bin/golint ./... $(GOPATH)/bin/golint ./...

View file

@ -5,6 +5,7 @@
:link: 'http://tlgrm.localhost/content' # webserver public address :link: 'http://tlgrm.localhost/content' # webserver public address
:upload: 'https:///xmppfiles.localhost' # xmpp http upload address :upload: 'https:///xmppfiles.localhost' # xmpp http upload address
:user: 'www-data' # owner of content files :user: 'www-data' # owner of content files
:quota: '256MB' # maximum storage size
:tdlib_verbosity: 1 :tdlib_verbosity: 1
:tdlib: :tdlib:
:datadir: './sessions/' :datadir: './sessions/'

View file

@ -39,6 +39,7 @@ type TelegramContentConfig struct {
Link string `yaml:":link"` Link string `yaml:":link"`
Upload string `yaml:":upload"` Upload string `yaml:":upload"`
User string `yaml:":user"` User string `yaml:":user"`
Quota string `yaml:":quota"`
} }
// TelegramTdlibConfig is for :tdlib: subtree // TelegramTdlibConfig is for :tdlib: subtree

View file

@ -24,6 +24,9 @@
}, },
":user": { ":user": {
"type": "string" "type": "string"
},
":quota": {
"type": "string"
} }
} }
}, },

View file

@ -361,6 +361,9 @@ func (c *Client) formatFile(file *client.File, compact bool) string {
return "" return ""
} }
gateway.StorageLock.Lock()
defer gateway.StorageLock.Unlock()
var link string var link string
var src string var src string
@ -372,6 +375,9 @@ func (c *Client) formatFile(file *client.File, compact bool) string {
return "" return ""
} }
size64:= uint64(file.Size)
c.prepareDiskSpace(size64)
basename := file.Remote.UniqueId + filepath.Ext(src) basename := file.Remote.UniqueId + filepath.Ext(src)
dest := c.content.Path + "/" + basename // destination path dest := c.content.Path + "/" + basename // destination path
link = c.content.Link + "/" + basename // download link link = c.content.Link + "/" + basename // download link
@ -387,6 +393,8 @@ func (c *Client) formatFile(file *client.File, compact bool) string {
return "<ERROR>" return "<ERROR>"
} }
} }
gateway.CachedStorageSize += size64
// chown // chown
if c.content.User != "" { if c.content.User != "" {
user, err := osUser.Lookup(c.content.User) user, err := osUser.Lookup(c.content.User)
@ -729,7 +737,12 @@ func (c *Client) messageToPrefix(message *client.Message, previewString string,
} }
func (c *Client) ensureDownloadFile(file *client.File) *client.File { func (c *Client) ensureDownloadFile(file *client.File) *client.File {
gateway.StorageLock.Lock()
defer gateway.StorageLock.Unlock()
if file != nil { if file != nil {
c.prepareDiskSpace(uint64(file.Size))
newFile, err := c.DownloadFile(file.Id, 1, true) newFile, err := c.DownloadFile(file.Id, 1, true)
if err == nil { if err == nil {
return newFile return newFile
@ -952,3 +965,16 @@ func (c *Client) subscribeToID(id int64, chat *client.Chat) {
args..., args...,
) )
} }
func (c *Client) prepareDiskSpace(size uint64) {
if gateway.StorageQuota > 0 && c.content.Path != "" {
var loweredQuota uint64
if gateway.StorageQuota >= size {
loweredQuota = gateway.StorageQuota - size
}
if gateway.CachedStorageSize >= loweredQuota {
log.Warn("Storage is rapidly clogged")
gateway.CleanOldFiles(c.content.Path, loweredQuota)
}
}
}

View file

@ -2,6 +2,8 @@ package xmpp
import ( import (
"github.com/pkg/errors" "github.com/pkg/errors"
"regexp"
"strconv"
"sync" "sync"
"time" "time"
@ -20,6 +22,19 @@ var sessions map[string]*telegram.Client
var db *persistence.SessionsYamlDB var db *persistence.SessionsYamlDB
var sessionLock sync.Mutex var sessionLock sync.Mutex
const (
B uint64 = 1
KB = B << 10
MB = KB << 10
GB = MB << 10
TB = GB << 10
PB = TB << 10
EB = PB << 10
maxUint64 uint64 = (1 << 64) - 1
)
var sizeRegex = regexp.MustCompile("\\A([0-9]+) ?([KMGTPE]?B?)\\z")
// NewComponent starts a new component and wraps it in // NewComponent starts a new component and wraps it in
// a stream manager that you should start yourself // a stream manager that you should start yourself
func NewComponent(conf config.XMPPConfig, tc config.TelegramConfig) (*xmpp.StreamManager, *xmpp.Component, error) { func NewComponent(conf config.XMPPConfig, tc config.TelegramConfig) (*xmpp.StreamManager, *xmpp.Component, error) {
@ -32,6 +47,13 @@ func NewComponent(conf config.XMPPConfig, tc config.TelegramConfig) (*xmpp.Strea
tgConf = tc tgConf = tc
if tc.Content.Quota != "" {
gateway.StorageQuota, err = parseSize(tc.Content.Quota)
if err != nil {
log.Warnf("Error parsing the storage quota: %v; the cleaner is disabled", err)
}
}
options := xmpp.ComponentOptions{ options := xmpp.ComponentOptions{
TransportConfiguration: xmpp.TransportConfiguration{ TransportConfiguration: xmpp.TransportConfiguration{
Address: conf.Host + ":" + conf.Port, Address: conf.Host + ":" + conf.Port,
@ -80,10 +102,22 @@ func heartbeat(component *xmpp.Component) {
} }
sessionLock.Unlock() sessionLock.Unlock()
quotaLowThreshold := gateway.StorageQuota / 10 * 9
log.Info("Starting heartbeat queue") log.Info("Starting heartbeat queue")
// status updater thread // status updater thread
for { for {
gateway.StorageLock.Lock()
if quotaLowThreshold > 0 && tgConf.Content.Path != "" {
gateway.MeasureStorageSize(tgConf.Content.Path)
if gateway.CachedStorageSize > quotaLowThreshold {
gateway.CleanOldFiles(tgConf.Content.Path, quotaLowThreshold)
}
}
gateway.StorageLock.Unlock()
time.Sleep(60e9) time.Sleep(60e9)
now := time.Now().Unix() now := time.Now().Unix()
@ -201,3 +235,47 @@ func Close(component *xmpp.Component) {
// close stream // close stream
component.Disconnect() component.Disconnect()
} }
// based on https://github.com/c2h5oh/datasize/blob/master/datasize.go
func parseSize(sSize string) (uint64, error) {
sizeParts := sizeRegex.FindStringSubmatch(sSize)
if len(sizeParts) > 2 {
numPart, err := strconv.ParseInt(sizeParts[1], 10, 64)
if err != nil {
return 0, err
}
var divisor uint64
val := uint64(numPart)
if len(sizeParts[2]) > 0 {
switch sizeParts[2][0] {
case 'B':
divisor = 1
case 'K':
divisor = KB
case 'M':
divisor = MB
case 'G':
divisor = GB
case 'T':
divisor = TB
case 'P':
divisor = PB
case 'E':
divisor = EB
}
}
if divisor == 0 {
return 0, &strconv.NumError{"Wrong suffix", sSize, strconv.ErrSyntax}
}
if val > maxUint64/divisor {
return 0, &strconv.NumError{"Overflow", sSize, strconv.ErrRange}
}
return val * divisor, nil
}
return 0, &strconv.NumError{"Not enough parts", sSize, strconv.ErrSyntax}
}

47
xmpp/component_test.go Normal file
View file

@ -0,0 +1,47 @@
package xmpp
import (
"testing"
)
func TestParseSizeGarbage(t *testing.T) {
_, err := parseSize("abc")
if err == nil {
t.Error("abc should not be accepted")
}
}
func TestParseSizeAsphalt(t *testing.T) {
size, err := parseSize("2B")
if size != 2 {
t.Errorf("Error parsing two bytes: %v %v", size, err)
}
}
func TestParseSize9K(t *testing.T) {
size, err := parseSize("9 KB")
if size != 9216 {
t.Errorf("Error parsing 9K: %v %v", size, err)
}
}
func TestParseSizeBits(t *testing.T) {
size, err := parseSize("9 Kb")
if err == nil {
t.Errorf("Error parsing kilobits: %v %v", size, err)
}
}
func TestParseSizeEB(t *testing.T) {
size, err := parseSize("3EB")
if size != 3458764513820540928 {
t.Errorf("Error parsing exabytes: %v %v", size, err)
}
}
func TestParseSizeOverflow(t *testing.T) {
size, err := parseSize("314EB")
if err == nil {
t.Errorf("Overflow is not overflowing: %v %v", size, err)
}
}

89
xmpp/gateway/storage.go Normal file
View file

@ -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
}
}
}