From 463be3b969a6d44de9ec5b22b4610d9f3d11cdbc Mon Sep 17 00:00:00 2001
From: stephentuso <tusodev@gmail.com>
Date: Thu, 4 Feb 2021 03:05:53 -0500
Subject: [PATCH] fix: add cache for file sizes

---
 .../companion/src/server/controllers/get.js   |  6 +-
 .../companion/src/server/provider/index.js    |  4 +-
 .../src/server/provider/withCache.js          | 60 +++++++++++++++++++
 3 files changed, 65 insertions(+), 5 deletions(-)
 create mode 100644 packages/@uppy/companion/src/server/provider/withCache.js

diff --git a/packages/@uppy/companion/src/server/controllers/get.js b/packages/@uppy/companion/src/server/controllers/get.js
index 0e6fef0115..5f341be3fe 100644
--- a/packages/@uppy/companion/src/server/controllers/get.js
+++ b/packages/@uppy/companion/src/server/controllers/get.js
@@ -5,8 +5,6 @@ const { retryWithDelay } = require('../helpers/utils')
 
 const workerCount = process.env.COMPANION_WORKER_COUNT ? parseInt(process.env.COMPANION_WORKER_COUNT, 10) : 3
 
-const retryDelays = [5000, 10000, 15000, 30000, 60000, 120000]
-
 const queue = (() => {
   const state = {}
 
@@ -56,7 +54,7 @@ function get (req, res, next) {
   const tenantId = req.body.metadata && req.body.metadata.tenantId
 
   retryWithDelay({
-    retryDelays,
+    retryDelays: [5000, 10000],
     action: () => new Promise((resolve, reject) => provider.size({ id, token, query: req.query }, (err, size) => {
       if (err) {
         reject(err)
@@ -115,7 +113,7 @@ function get (req, res, next) {
       uploader.onSocketReady(() => {
         if (tenantId) {
           queue(tenantId, () => retryWithDelay({
-            retryDelays,
+            retryDelays: [5000, 10000, 15000, 30000, 60000, 120000],
             action: getPerformDownload(false),
             lastAction: getPerformDownload(true)
           }))
diff --git a/packages/@uppy/companion/src/server/provider/index.js b/packages/@uppy/companion/src/server/provider/index.js
index 2f81254b37..08fe4beb86 100644
--- a/packages/@uppy/companion/src/server/provider/index.js
+++ b/packages/@uppy/companion/src/server/provider/index.js
@@ -18,6 +18,7 @@ const { getCredentialsResolver } = require('./credentials')
 const Provider = require('./Provider')
 // eslint-disable-next-line
 const SearchProvider = require('./SearchProvider')
+const { withCache } = require('./withCache')
 
 // leave here for now until Purest Providers gets updated with Zoom provider
 config.zoom = {
@@ -61,7 +62,8 @@ module.exports.getProviderMiddleware = (providers, needsProviderCredentials) =>
    */
   const middleware = (req, res, next, providerName) => {
     if (providers[providerName] && validOptions(req.companion.options)) {
-      req.companion.provider = new providers[providerName]({ providerName, config })
+      const provider = new providers[providerName]({ providerName, config })
+      req.companion.provider = provider instanceof Provider ? withCache(provider) : provider
       if (needsProviderCredentials) {
         req.companion.getProviderCredentials = getCredentialsResolver(providerName, req.companion.options, req)
       }
diff --git a/packages/@uppy/companion/src/server/provider/withCache.js b/packages/@uppy/companion/src/server/provider/withCache.js
new file mode 100644
index 0000000000..acfaf90354
--- /dev/null
+++ b/packages/@uppy/companion/src/server/provider/withCache.js
@@ -0,0 +1,60 @@
+const cache = {}
+
+const sizeKey = (providerName, id) => `${providerName}${id}size`
+
+const addToCache = (key, data) => {
+  cache[key] = data
+}
+
+const getFromCache = (key) => cache[key]
+
+/**
+ *
+ * @param {*} provider
+ */
+module.exports.withCache = (provider) => {
+  const providerName = provider.constructor.authProvider
+
+  const wrapList = (list) => (options, cb) => {
+    list(options, (err, data) => {
+      if (err) {
+        cb(err, data)
+        return
+      }
+
+      for (const item of data.items) {
+        if (item.size == null || item.id == null) {
+          continue
+        }
+        addToCache(sizeKey(providerName, item.id), item.size)
+      }
+
+      cb(err, data)
+    })
+  }
+
+  const wrapSize = (size) => (options, cb) => {
+    const key = sizeKey(providerName, options.id)
+    const cached = getFromCache(key)
+    if (cached) {
+      cb(null, cached)
+    }
+
+    size(options, (err, size) => {
+      if (err) {
+        cb(err, size)
+        return
+      }
+
+      addToCache(key, size)
+
+      cb(err, size)
+    })
+  }
+
+  // modifying is ugly but whatever
+  provider.list = wrapList(provider.list.bind(provider))
+  provider.size = wrapSize(provider.size.bind(provider))
+
+  return provider
+}