Optimizations

2023-11-28 22:57:25 +00:00
parent e7e5806b20
commit e797824ab0
8 changed files with 139 additions and 101 deletions
--- a/internal/dav/listing.go
+++ b/internal/dav/listing.go
@@ -9,21 +9,24 @@ import (

 	"github.com/debridmediamanager.com/zurg/internal/torrent"
 	"github.com/debridmediamanager.com/zurg/pkg/dav"
+	"github.com/dgraph-io/ristretto"
 	"go.uber.org/zap"
 )

-func HandlePropfindRequest(w http.ResponseWriter, r *http.Request, t *torrent.TorrentManager, log *zap.SugaredLogger) {
+func HandlePropfindRequest(w http.ResponseWriter, r *http.Request, t *torrent.TorrentManager, cache *ristretto.Cache, log *zap.SugaredLogger) {
 	requestPath := path.Clean(r.URL.Path)
-	filteredSegments := splitIntoSegments(requestPath)

+	var output *string
 	var err error
+
+	filteredSegments := splitIntoSegments(requestPath)
 	switch {
 	case len(filteredSegments) == 0:
-		err = handleListDirectories(w, t)
+		output, err = handleListDirectories(w, t)
 	case len(filteredSegments) == 1:
-		err = handleListTorrents(w, requestPath, t)
+		output, err = handleListTorrents(w, requestPath, t)
 	case len(filteredSegments) == 2:
-		err = handleListFiles(w, requestPath, t)
+		output, err = handleListFiles(w, requestPath, t)
 	default:
 		log.Warnf("Request %s %s not found", r.Method, requestPath)
 		http.Error(w, "Not Found", http.StatusNotFound)
@@ -40,13 +43,17 @@ func HandlePropfindRequest(w http.ResponseWriter, r *http.Request, t *torrent.To
 		return
 	}

-	w.Header().Set("Content-Type", "text/xml; charset=\"utf-8\"")
+	if output != nil {
+		w.Header().Set("Content-Type", "text/xml; charset=\"utf-8\"")
+		w.WriteHeader(http.StatusOK)
+		fmt.Fprint(w, *output)
+	}
 }

-func handleListDirectories(w http.ResponseWriter, t *torrent.TorrentManager) error {
-	fmt.Fprint(w, "<?xml version=\"1.0\" encoding=\"utf-8\"?><d:multistatus xmlns:d=\"DAV:\">")
+func handleListDirectories(w http.ResponseWriter, t *torrent.TorrentManager) (*string, error) {
+	davDoc := "<?xml version=\"1.0\" encoding=\"utf-8\"?><d:multistatus xmlns:d=\"DAV:\">"
 	// initial response is the directory itself
-	fmt.Fprint(w, dav.BaseDirectory("", ""))
+	davDoc += dav.BaseDirectory("", "")

 	directories := t.DirectoryMap.Keys()
 	sort.Strings(directories)
@@ -54,61 +61,40 @@ func handleListDirectories(w http.ResponseWriter, t *torrent.TorrentManager) err
 		if strings.HasPrefix(directory, "int__") {
 			continue
 		}
-		fmt.Fprint(w, dav.Directory(directory, ""))
+		davDoc += dav.Directory(directory, "")
 	}

-	fmt.Fprint(w, "</d:multistatus>")
-	return nil
+	davDoc += "</d:multistatus>"
+	return &davDoc, nil
 }

-func handleListTorrents(w http.ResponseWriter, requestPath string, t *torrent.TorrentManager) error {
+func handleListTorrents(w http.ResponseWriter, requestPath string, t *torrent.TorrentManager) (*string, error) {
 	basePath := path.Base(requestPath)
-	torrents, ok := t.DirectoryMap.Get(basePath)
+	_, ok := t.DirectoryMap.Get(basePath)
 	if !ok {
-		return fmt.Errorf("cannot find directory %s", basePath)
+		return nil, fmt.Errorf("cannot find directory %s", basePath)
 	}

-	fmt.Fprint(w, "<?xml version=\"1.0\" encoding=\"utf-8\"?><d:multistatus xmlns:d=\"DAV:\">")
-	// initial response is the directory itself
-	fmt.Fprint(w, dav.BaseDirectory(basePath, ""))
+	resp, _ := t.ResponseCache.Get(basePath + ".dav")
+	davDoc := "<?xml version=\"1.0\" encoding=\"utf-8\"?><d:multistatus xmlns:d=\"DAV:\">" + dav.BaseDirectory(basePath, "") + dav.BaseDirectory(basePath, "") + resp.(string) + "</d:multistatus>"

-	var allTorrents []torrent.Torrent
-	torrents.IterCb(func(key string, tor *torrent.Torrent) {
-		if tor.AllInProgress() {
-			return
-		}
-		copy := *tor
-		copy.AccessKey = key
-		allTorrents = append(allTorrents, copy)
-	})
-	sort.Slice(allTorrents, func(i, j int) bool {
-		return allTorrents[i].AccessKey < allTorrents[j].AccessKey
-	})
-
-	for _, tor := range allTorrents {
-		fmt.Fprint(w, dav.Directory(tor.AccessKey, tor.LatestAdded))
-	}
-
-	fmt.Fprint(w, "</d:multistatus>")
-	return nil
+	return &davDoc, nil
 }

-func handleListFiles(w http.ResponseWriter, requestPath string, t *torrent.TorrentManager) error {
+func handleListFiles(w http.ResponseWriter, requestPath string, t *torrent.TorrentManager) (*string, error) {
 	requestPath = strings.Trim(requestPath, "/")
 	basePath := path.Base(path.Dir(requestPath))
 	torrents, ok := t.DirectoryMap.Get(basePath)
 	if !ok {
-		return fmt.Errorf("cannot find directory %s", basePath)
+		return nil, fmt.Errorf("cannot find directory %s", basePath)
 	}
 	accessKey := path.Base(requestPath)
 	tor, ok := torrents.Get(accessKey)
 	if !ok {
-		return fmt.Errorf("cannot find torrent %s", accessKey)
+		return nil, fmt.Errorf("cannot find torrent %s", accessKey)
 	}

-	fmt.Fprint(w, "<?xml version=\"1.0\" encoding=\"utf-8\"?><d:multistatus xmlns:d=\"DAV:\">")
-	// initial response is the directory itself
-	fmt.Fprint(w, dav.BaseDirectory(requestPath, tor.LatestAdded))
+	davDoc := "<?xml version=\"1.0\" encoding=\"utf-8\"?><d:multistatus xmlns:d=\"DAV:\">" + dav.BaseDirectory(requestPath, tor.LatestAdded)

 	filenames := tor.SelectedFiles.Keys()
 	sort.Strings(filenames)
@@ -117,9 +103,9 @@ func handleListFiles(w http.ResponseWriter, requestPath string, t *torrent.Torre
 		if file == nil || !strings.HasPrefix(file.Link, "http") {
 			continue
 		}
-		fmt.Fprint(w, dav.File(filename, file.Bytes, file.Ended))
+		davDoc += dav.File(filename, file.Bytes, file.Ended)
 	}

-	fmt.Fprint(w, "</d:multistatus>")
-	return nil
+	davDoc += "</d:multistatus>"
+	return &davDoc, nil
 }
--- a/internal/http/listing.go
+++ b/internal/http/listing.go
@@ -10,10 +10,11 @@ import (
 	"strings"

 	"github.com/debridmediamanager.com/zurg/internal/torrent"
+	"github.com/dgraph-io/ristretto"
 	"go.uber.org/zap"
 )

-func HandleDirectoryListing(w http.ResponseWriter, r *http.Request, t *torrent.TorrentManager, log *zap.SugaredLogger) {
+func HandleDirectoryListing(w http.ResponseWriter, r *http.Request, t *torrent.TorrentManager, cache *ristretto.Cache, log *zap.SugaredLogger) {
 	requestPath := path.Clean(r.URL.Path)

 	var output *string
@@ -66,29 +67,14 @@ func handleRoot(t *torrent.TorrentManager) (*string, error) {

 func handleListOfTorrents(requestPath string, t *torrent.TorrentManager) (*string, error) {
 	basePath := path.Base(requestPath)
-	torrents, ok := t.DirectoryMap.Get(basePath)
+	_, ok := t.DirectoryMap.Get(basePath)
 	if !ok {
 		return nil, fmt.Errorf("cannot find directory %s", basePath)
 	}

-	htmlDoc := "<ol>"
+	resp, _ := t.ResponseCache.Get(basePath + ".html")
+	htmlDoc := resp.(string)

-	var allTorrents []torrent.Torrent
-	torrents.IterCb(func(key string, tor *torrent.Torrent) {
-		if tor.AllInProgress() {
-			return
-		}
-		copy := *tor
-		copy.AccessKey = key
-		allTorrents = append(allTorrents, copy)
-	})
-	sort.Slice(allTorrents, func(i, j int) bool {
-		return allTorrents[i].AccessKey < allTorrents[j].AccessKey
-	})
-
-	for _, tor := range allTorrents {
-		htmlDoc = htmlDoc + fmt.Sprintf("<li><a href=\"%s/\">%s</a></li>", filepath.Join(requestPath, url.PathEscape(tor.AccessKey)), tor.AccessKey)
-	}
 	return &htmlDoc, nil
 }

--- a/internal/net/router.go
+++ b/internal/net/router.go
@@ -10,19 +10,20 @@ import (
 	intHttp "github.com/debridmediamanager.com/zurg/internal/http"
 	"github.com/debridmediamanager.com/zurg/internal/torrent"
 	"github.com/debridmediamanager.com/zurg/internal/universal"
+	"github.com/dgraph-io/ristretto"
 	"go.uber.org/zap"
 )

 // Router creates a WebDAV router
-func Router(mux *http.ServeMux, getfile *universal.GetFile, c config.ConfigInterface, t *torrent.TorrentManager, log *zap.SugaredLogger) {
+func Router(mux *http.ServeMux, getfile *universal.GetFile, c config.ConfigInterface, t *torrent.TorrentManager, cache *ristretto.Cache, log *zap.SugaredLogger) {
 	mux.HandleFunc("/http/", func(w http.ResponseWriter, r *http.Request) {
 		switch r.Method {
 		case http.MethodGet:
 			requestPath := path.Clean(r.URL.Path)
 			if countNonEmptySegments(strings.Split(requestPath, "/")) > 3 {
-				getfile.HandleGetRequest(w, r, t, c, log)
+				getfile.HandleGetRequest(w, r, t, c, cache, log)
 			} else {
-				intHttp.HandleDirectoryListing(w, r, t, log)
+				intHttp.HandleDirectoryListing(w, r, t, cache, log)
 			}

 		case http.MethodHead:
@@ -37,13 +38,13 @@ func Router(mux *http.ServeMux, getfile *universal.GetFile, c config.ConfigInter
 	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
 		switch r.Method {
 		case "PROPFIND":
-			dav.HandlePropfindRequest(w, r, t, log)
+			dav.HandlePropfindRequest(w, r, t, cache, log)

 		case http.MethodDelete:
 			dav.HandleDeleteRequest(w, r, t, log)

 		case http.MethodGet:
-			getfile.HandleGetRequest(w, r, t, c, log)
+			getfile.HandleGetRequest(w, r, t, c, cache, log)

 		case http.MethodOptions:
 			w.WriteHeader(http.StatusOK)
--- a/internal/torrent/manager.go
+++ b/internal/torrent/manager.go
@@ -7,13 +7,16 @@ import (
 	"math"
 	"os"
 	"path/filepath"
+	"sort"
 	"strings"
 	"sync"
 	"time"

 	"github.com/debridmediamanager.com/zurg/internal/config"
+	"github.com/debridmediamanager.com/zurg/pkg/dav"
 	"github.com/debridmediamanager.com/zurg/pkg/realdebrid"
 	"github.com/debridmediamanager.com/zurg/pkg/utils"
+	"github.com/dgraph-io/ristretto"
 	cmap "github.com/orcaman/concurrent-map/v2"
 	"github.com/panjf2000/ants/v2"
 	"go.uber.org/zap"
@@ -26,13 +29,14 @@ const (
 )

 type TorrentManager struct {
+	Config          config.ConfigInterface
+	Api             *realdebrid.RealDebrid
 	DirectoryMap    cmap.ConcurrentMap[string, cmap.ConcurrentMap[string, *Torrent]] // directory -> accessKey -> Torrent
 	DownloadCache   cmap.ConcurrentMap[string, *realdebrid.Download]
+	ResponseCache   *ristretto.Cache
 	checksum        string
 	latestAdded     string
 	requiredVersion string
-	cfg             config.ConfigInterface
-	Api             *realdebrid.RealDebrid
 	antsPool        *ants.Pool
 	unrestrictPool  *ants.Pool
 	log             *zap.SugaredLogger
@@ -42,18 +46,19 @@ type TorrentManager struct {
 // NewTorrentManager creates a new torrent manager
 // it will fetch all torrents and their info in the background
 // and store them in-memory and cached in files
-func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, p *ants.Pool, log *zap.SugaredLogger) *TorrentManager {
+func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, p *ants.Pool, cache *ristretto.Cache, log *zap.SugaredLogger) *TorrentManager {
 	t := &TorrentManager{
-		cfg:             cfg,
-		DirectoryMap:    cmap.New[cmap.ConcurrentMap[string, *Torrent]](),
-		requiredVersion: "18.11.2023",
+		Config:          cfg,
 		Api:             api,
+		DirectoryMap:    cmap.New[cmap.ConcurrentMap[string, *Torrent]](),
+		ResponseCache:   cache,
+		requiredVersion: "18.11.2023",
 		antsPool:        p,
 		log:             log,
 		mu:              &sync.Mutex{},
 	}

-	unrestrictPool, err := ants.NewPool(t.cfg.GetUnrestrictWorkers())
+	unrestrictPool, err := ants.NewPool(t.Config.GetUnrestrictWorkers())
 	if err != nil {
 		t.unrestrictPool = t.antsPool
 	} else {
@@ -145,12 +150,12 @@ func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, p
 		// get filenames
 		filenames := torrent.SelectedFiles.Keys()
 		// Map torrents to directories
-		switch t.cfg.GetVersion() {
+		switch t.Config.GetVersion() {
 		case "v1":
-			configV1 := t.cfg.(*config.ZurgConfigV1)
+			configV1 := t.Config.(*config.ZurgConfigV1)
 			for _, directories := range configV1.GetGroupMap() {
 				for _, directory := range directories {
-					if t.cfg.MeetsConditions(directory, torrent.AccessKey, torrentIDs, filenames) {
+					if t.Config.MeetsConditions(directory, torrent.AccessKey, torrentIDs, filenames) {
 						torrents, _ := t.DirectoryMap.Get(directory)
 						torrents.Set(accessKey, torrent)
 						break
@@ -159,12 +164,13 @@ func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, p
 			}
 		}
 	})
+	t.updateSortedKeys()

 	t.log.Infof("Compiled into %d torrents, %d were missing info", allTorrents.Count(), noInfoCount)

 	t.SetChecksum(t.getChecksum())

-	if t.cfg.EnableRepair() {
+	if t.Config.EnableRepair() {
 		t.log.Info("Checking for torrents to repair")
 		t.repairAll()
 		t.log.Info("Finished checking for torrents to repair")
@@ -209,8 +215,8 @@ func (t *TorrentManager) mergeToMain(mainTorrent, torrentToMerge *Torrent) *Torr
 func (t *TorrentManager) UnrestrictUntilOk(link string) *realdebrid.Download {
 	retChan := make(chan *realdebrid.Download, 1)
 	t.unrestrictPool.Submit(func() {
-		retChan <- t.Api.UnrestrictUntilOk(link, t.cfg.ShouldServeFromRclone())
-		time.Sleep(time.Duration(t.cfg.GetReleaseUnrestrictAfterMs()) * time.Millisecond)
+		retChan <- t.Api.UnrestrictUntilOk(link, t.Config.ShouldServeFromRclone())
+		time.Sleep(time.Duration(t.Config.GetReleaseUnrestrictAfterMs()) * time.Millisecond)
 	})
 	defer close(retChan)
 	return <-retChan
@@ -282,7 +288,7 @@ func (t *TorrentManager) getChecksum() string {
 func (t *TorrentManager) startRefreshJob() {
 	t.log.Info("Starting periodic refresh")
 	for {
-		<-time.After(time.Duration(t.cfg.GetRefreshEverySeconds()) * time.Second)
+		<-time.After(time.Duration(t.Config.GetRefreshEverySeconds()) * time.Second)

 		checksum := t.getChecksum()
 		if checksum == t.checksum {
@@ -358,12 +364,12 @@ func (t *TorrentManager) startRefreshJob() {
 			// get filenames
 			filenames := torrent.SelectedFiles.Keys()
 			// Map torrents to directories
-			switch t.cfg.GetVersion() {
+			switch t.Config.GetVersion() {
 			case "v1":
-				configV1 := t.cfg.(*config.ZurgConfigV1)
+				configV1 := t.Config.(*config.ZurgConfigV1)
 				for _, directories := range configV1.GetGroupMap() {
 					for _, directory := range directories {
-						if t.cfg.MeetsConditions(directory, torrent.AccessKey, torrentIDs, filenames) {
+						if t.Config.MeetsConditions(directory, torrent.AccessKey, torrentIDs, filenames) {
 							torrents, _ := t.DirectoryMap.Get(directory)
 							torrents.Set(torrent.AccessKey, torrent)
 							if torrent.LatestAdded > t.latestAdded {
@@ -386,19 +392,20 @@ func (t *TorrentManager) startRefreshJob() {
 				t.log.Infof("Deleted torrent: %s\n", oldAccessKey)
 			}
 		}
+		t.updateSortedKeys()

 		t.log.Infof("Compiled into %d torrents, %d were missing info", oldTorrents.Count(), noInfoCount)

 		t.SetChecksum(t.getChecksum())

-		if t.cfg.EnableRepair() {
+		if t.Config.EnableRepair() {
 			t.log.Info("Checking for torrents to repair")
 			t.repairAll()
 			t.log.Info("Finished checking for torrents to repair")
 		} else {
 			t.log.Info("Repair is disabled, skipping repair check")
 		}
-		go OnLibraryUpdateHook(updatedPaths, t.cfg, t.log)
+		go OnLibraryUpdateHook(updatedPaths, t.Config, t.log)

 		t.latestAdded = newTorrents[0].Added
 		t.log.Info("Finished refreshing torrents")
@@ -477,11 +484,11 @@ func (t *TorrentManager) getMoreInfo(rdTorrent realdebrid.Torrent) *Torrent {
 }

 func (t *TorrentManager) getName(name, originalName string) string {
-	if t.cfg.EnableRetainRDTorrentName() {
+	if t.Config.EnableRetainRDTorrentName() {
 		return name
 	}
 	// drop the extension from the name
-	if t.cfg.EnableRetainFolderNameExtension() && strings.Contains(name, originalName) {
+	if t.Config.EnableRetainFolderNameExtension() && strings.Contains(name, originalName) {
 		return name
 	} else {
 		ret := strings.TrimSuffix(originalName, ".mp4")
@@ -551,9 +558,9 @@ func (t *TorrentManager) organizeChaos(links []string, selectedFiles []*File) ([
 				resultsChan <- Result{Response: download}
 				return
 			}
-			resp := t.Api.UnrestrictUntilOk(link, t.cfg.ShouldServeFromRclone())
+			resp := t.Api.UnrestrictUntilOk(link, t.Config.ShouldServeFromRclone())
 			resultsChan <- Result{Response: resp}
-			time.Sleep(time.Duration(t.cfg.GetReleaseUnrestrictAfterMs()) * time.Millisecond)
+			time.Sleep(time.Duration(t.Config.GetReleaseUnrestrictAfterMs()) * time.Millisecond)
 		})
 	}

@@ -649,10 +656,11 @@ func (t *TorrentManager) Delete(accessKey string) {
 			torrents.Remove(accessKey)
 		}
 	})
+	t.updateSortedKeys()
 }

 func (t *TorrentManager) Repair(accessKey string) {
-	if !t.cfg.EnableRepair() {
+	if !t.Config.EnableRepair() {
 		t.log.Warn("Repair is disabled; if you do not have other zurg instances running, you should enable repair")
 		return
 	}
@@ -700,6 +708,7 @@ func (t *TorrentManager) Repair(accessKey string) {
 		t.DirectoryMap.IterCb(func(_ string, torrents cmap.ConcurrentMap[string, *Torrent]) {
 			torrents.Remove(torrent.AccessKey)
 		})
+		t.updateSortedKeys()
 		// t.log.Debugf("You can try fixing it yourself magnet:?xt=urn:btih:%s", info.Hash)
 		return
 	} else if streamableCount == 1 {
@@ -708,6 +717,7 @@ func (t *TorrentManager) Repair(accessKey string) {
 		t.DirectoryMap.IterCb(func(_ string, torrents cmap.ConcurrentMap[string, *Torrent]) {
 			torrents.Remove(torrent.AccessKey)
 		})
+		t.updateSortedKeys()
 		return
 	}
 	// t.log.Debugf("Identified the expired files of torrent id=%s", info.ID)
@@ -876,3 +886,23 @@ func (t *TorrentManager) canCapacityHandle() bool {
 		retryCount++
 	}
 }
+
+func (t *TorrentManager) updateSortedKeys() {
+	t.DirectoryMap.IterCb(func(directory string, torrents cmap.ConcurrentMap[string, *Torrent]) {
+		allKeys := torrents.Keys()
+		sort.Strings(allKeys)
+		davRet := ""
+		htmlRet := ""
+		for _, accessKey := range allKeys {
+			if tor, ok := torrents.Get(accessKey); ok {
+				if tor.AnyInProgress() {
+					continue
+				}
+				davRet += dav.Directory(tor.AccessKey, tor.LatestAdded)
+				htmlRet += fmt.Sprintf("<a href=\"/%s/%s\">%s</a><br>", directory, tor.AccessKey, tor.AccessKey)
+			}
+		}
+		t.ResponseCache.Set(directory+".dav", davRet, 1)
+		t.ResponseCache.Set(directory+".html", "<ol>"+htmlRet, 1)
+	})
+}
--- a/internal/universal/get.go
+++ b/internal/universal/get.go
@@ -14,6 +14,7 @@ import (
 	intHttp "github.com/debridmediamanager.com/zurg/internal/http"
 	intTor "github.com/debridmediamanager.com/zurg/internal/torrent"
 	zurghttp "github.com/debridmediamanager.com/zurg/pkg/http"
+	"github.com/dgraph-io/ristretto"
 	"go.uber.org/zap"
 )

@@ -26,7 +27,7 @@ func NewGetFile(client *zurghttp.HTTPClient) *GetFile {
 }

 // HandleGetRequest handles a GET request universally for both WebDAV and HTTP
-func (gf *GetFile) HandleGetRequest(w http.ResponseWriter, r *http.Request, t *intTor.TorrentManager, c config.ConfigInterface, log *zap.SugaredLogger) {
+func (gf *GetFile) HandleGetRequest(w http.ResponseWriter, r *http.Request, t *intTor.TorrentManager, c config.ConfigInterface, cache *ristretto.Cache, log *zap.SugaredLogger) {
 	requestPath := path.Clean(r.URL.Path)
 	isDav := true
 	if strings.Contains(requestPath, "/http") {
@@ -40,9 +41,9 @@ func (gf *GetFile) HandleGetRequest(w http.ResponseWriter, r *http.Request, t *i
 	// If there are less than 3 segments, return an error or adjust as needed
 	if len(segments) <= 3 {
 		if isDav {
-			dav.HandlePropfindRequest(w, r, t, log)
+			dav.HandlePropfindRequest(w, r, t, cache, log)
 		} else {
-			intHttp.HandleDirectoryListing(w, r, t, log)
+			intHttp.HandleDirectoryListing(w, r, t, cache, log)
 		}
 		return
 	}