Add caching of urls (#38)

* added cache.py as part of modules

* added argument for cacheSize and made metrics for cache

* refactor /find endpoint to use new cache

* fixed up a bit of commits

* Update modules/cache.py

Co-authored-by: pnavab <114110926+pnavab@users.noreply.github.com>

* Update modules/cache.py

Co-authored-by: pnavab <114110926+pnavab@users.noreply.github.com>

* reworded and reformated code for better readability

* minor fixes

* more small changes

* remove MetricsHandler instance method, cache size uses gauge

* remove call to instance

* linting

---------

Co-authored-by: pnavab <114110926+pnavab@users.noreply.github.com>
Co-authored-by: evan ugarte <evanuxd@gmail.com>

Author: 3 people

modules/args.py

@@ -32,5 +32,11 @@ def get_args():
         default=0,
         help="increase logging verbosity; can be used multiple times",
     )
+    parser.add_argument(
+        "--cache-size",
+        type=int,
+        default=100,
+        help="number of url redirects to store in memory. defaults to 100"
+    )
 
     return parser.parse_args()

modules/cache.py

@@ -0,0 +1,27 @@
+from collections import OrderedDict
+import logging
+
+from modules.metrics import MetricsHandler
+
+class Cache:
+    def __init__(self, cacheSize):
+        self.dict = OrderedDict()
+        self.size = cacheSize
+
+    def find(self, alias):
+        if alias not in self.dict:
+            MetricsHandler.cache_misses.inc()
+            return None
+        
+        self.dict.move_to_end(alias, last=False) #move alias to front of cache
+        logging.debug(f"alias: '{alias}' is grabbed from mapping")
+        MetricsHandler.cache_hits.inc()
+        return self.dict[alias]
+
+    def add(self, alias, url_output):
+        if len(self.dict) == self.size: 
+            data = self.dict.popitem() #remove least used alias if size reaches max
+            logging.debug(f"alias: {data[0]} has been removed from cache")
+        self.dict[alias] = url_output
+        MetricsHandler.cache_size.set(len(self.dict))
+        logging.debug("set alias: '" + alias + "' to mapping")

modules/metrics.py

@@ -21,6 +21,21 @@ class Metrics(enum.Enum):
         prometheus_client.Counter,
         ['path', 'code'],
     )
+    CACHE_SIZE = (
+        "cache_size",
+        "Size of LRU cache for /find",
+        prometheus_client.Gauge,
+    ) 
+    CACHE_HITS = (
+        "cache_hits",
+        "Number of times cache is used",
+        prometheus_client.Counter,
+    )
+    CACHE_MIESSES = (
+        "cache_misses",
+        "Number of times cahes is not used",
+        prometheus_client.Counter,
+    )
 
     def __init__(self, title, description, prometheus_type, labels=()):
         # we use the above default value for labels because it matches what's used
@@ -33,11 +48,7 @@ def __init__(self, title, description, prometheus_type, labels=()):
 
 
 class MetricsHandler:
-    _instance = None
-
-    def __init__(self):
-        raise RuntimeError('Call MetricsHandler.instance() instead')
-
+    @classmethod
     def init(self) -> None:
         for metric in Metrics:
             setattr(
@@ -47,10 +58,3 @@ def init(self) -> None:
                     metric.title, metric.description, labelnames=metric.labels
                 ),
             )
-
-    @classmethod
-    def instance(cls):
-        if cls._instance is None:
-            cls._instance = cls.__new__(cls)
-            cls.init(cls)
-        return cls._instance

server.py

@@ -15,6 +15,7 @@
 from modules.constants import HttpResponse, http_code_to_enum
 from modules.metrics import MetricsHandler
 from modules.sqlite_helpers import increment_used_column
+from modules.cache import Cache
 
 
 app = FastAPI()
@@ -23,17 +24,18 @@
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=['*'],
-    allow_methods=['*'],
-    allow_headers=['*'],
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
 )
 
-metrics_handler = MetricsHandler.instance()
+cache = Cache(args.cache_size)
 
-#maybe create the table if it doesnt already exist
+# maybe create the table if it doesnt already exist
 DATABASE_FILE = args.database_file_path
 sqlite_helpers.maybe_create_table(DATABASE_FILE)
 
+
 # middleware to get metrics on HTTP response codes
 @app.middleware("http")
 async def track_response_codes(request: Request, call_next):
@@ -42,38 +44,45 @@ async def track_response_codes(request: Request, call_next):
     MetricsHandler.http_code.labels(request.url.path, status_code).inc()
     return response
 
+
 @app.post("/create_url")
 async def create_url(request: Request):
     urljson = await request.json()
     logging.debug(f"/create_url called with body: {urljson}")
     alias = None
 
     try:
-        alias = urljson.get('alias')
+        alias = urljson.get("alias")
         if alias is None:
             if args.disable_random_alias:
                 raise KeyError("alias must be specified")
             else:
-                alias = generate_alias(urljson['url'])
+                alias = generate_alias(urljson["url"])
         if not alias.isalnum():
             raise ValueError("alias must only contain alphanumeric characters")
 
         with MetricsHandler.query_time.labels("create").time():
-            response = sqlite_helpers.insert_url(DATABASE_FILE, urljson['url'], alias)
+            response = sqlite_helpers.insert_url(DATABASE_FILE, urljson["url"], alias)
             if response is not None:
                 MetricsHandler.url_count.inc(1)
-                return { "url": urljson['url'], "alias": alias, "created_at": response}
+                return {"url": urljson["url"], "alias": alias, "created_at": response}
             else:
-                raise HTTPException(status_code=HttpResponse.CONFLICT.code )
+                raise HTTPException(status_code=HttpResponse.CONFLICT.code)
     except KeyError:
         logging.exception("returning 400 due to missing key")
         raise HTTPException(status_code=HttpResponse.BAD_REQUEST.code)
     except ValueError:
-        logging.exception(f"returning 422 due to invalid alias of \"{alias}\"")
+        logging.exception(f'returning 422 due to invalid alias of "{alias}"')
         raise HTTPException(status_code=HttpResponse.INVALID_ARGUMENT_EXCEPTION.code)
 
+
 @app.get("/list")
-async def get_urls(search: Optional[str] = None, page: int = 0, sort_by: str = "created_at", order: str = "DESC"):
+async def get_urls(
+    search: Optional[str] = None,
+    page: int = 0,
+    sort_by: str = "created_at",
+    order: str = "DESC",
+):
     valid_sort_attributes = {"id", "url", "alias", "created_at", "used"}
     if order not in {"DESC", "ASC"}:
         raise HTTPException(status_code=400, detail="Invalid order")
@@ -82,20 +91,36 @@ async def get_urls(search: Optional[str] = None, page: int = 0, sort_by: str = "
     if page < 0:
         raise HTTPException(status_code=400, detail="Invalid page number")
     if search and not search.isalnum():
-        raise HTTPException(status_code=400, detail=f'search term "{search}" is invalid. only alphanumeric chars are allowed')
+        raise HTTPException(
+            status_code=400,
+            detail=f'search term "{search}" is invalid. only alphanumeric chars are allowed',
+        )
     with MetricsHandler.query_time.labels("list").time():
-        urls = sqlite_helpers.get_urls(DATABASE_FILE, page, search=search, sort_by=sort_by, order=order)
+        urls = sqlite_helpers.get_urls(
+            DATABASE_FILE, page, search=search, sort_by=sort_by, order=order
+        )
         total_urls = sqlite_helpers.get_number_of_entries(DATABASE_FILE, search=search)
-        return {"data": urls, "total": total_urls, "rows_per_page": sqlite_helpers.ROWS_PER_PAGE}
+        return {
+            "data": urls,
+            "total": total_urls,
+            "rows_per_page": sqlite_helpers.ROWS_PER_PAGE,
+        }
+
 
 @app.get("/find/{alias}")
 async def get_url(alias: str):
     logging.debug(f"/find called with alias: {alias}")
+    url_output = cache.find(alias)  # try to find url in cache
+    if url_output is not None:
+        alias_queue.put(alias)
+        return RedirectResponse(url_output)
+
     with MetricsHandler.query_time.labels("find").time():
         url_output = sqlite_helpers.get_url(DATABASE_FILE, alias)
-
     if url_output is None:
         raise HTTPException(status_code=HttpResponse.NOT_FOUND.code)
+    cache.add(alias, url_output)  # else, adds url and alias to cache
+
     alias_queue.put(alias)
     return RedirectResponse(url_output)
 
@@ -104,15 +129,19 @@ async def get_url(alias: str):
 async def delete_url(alias: str):
     logging.debug(f"/delete called with alias: {alias}")
     with MetricsHandler.query_time.labels("delete").time():
-      if(sqlite_helpers.delete_url(DATABASE_FILE, alias)):
-          return {"message": "URL deleted successfully"}
-      else:
-          raise HTTPException(status_code=HttpResponse.NOT_FOUND.code)
+        if sqlite_helpers.delete_url(DATABASE_FILE, alias):
+            return {"message": "URL deleted successfully"}
+        else:
+            raise HTTPException(status_code=HttpResponse.NOT_FOUND.code)
+
 
 @app.exception_handler(HTTPException)
 async def http_exception_handler(request, exc):
     status_code_enum = http_code_to_enum[exc.status_code]
-    return HTMLResponse(content=status_code_enum.content, status_code=status_code_enum.code)
+    return HTMLResponse(
+        content=status_code_enum.content, status_code=status_code_enum.code
+    )
+
 
 @app.get("/metrics")
 def get_metrics():
@@ -121,20 +150,22 @@ def get_metrics():
         content=prometheus_client.generate_latest(),
     )
 
+
 logging.Formatter.converter = time.gmtime
 
 logging.basicConfig(
     # in mondo we trust
     format="%(asctime)s.%(msecs)03dZ %(levelname)s:%(name)s:%(message)s",
     datefmt="%Y-%m-%dT%H:%M:%S",
-    level= logging.ERROR - (args.verbose*10),
+    level=logging.ERROR - (args.verbose * 10),
 )
 
+
 def consumer():
     while True:
         alias = alias_queue.get()
-        if alias is None:  
-            break  
+        if alias is None:
+            break
         try:
             with MetricsHandler.query_time.labels("increment_used").time():
                 increment_used_column(DATABASE_FILE, alias)
@@ -144,7 +175,6 @@ def consumer():
             alias_queue.task_done()
 
 
-    
 # we have a separate __name__ check here due to how FastAPI starts
 # a server. the file is first ran (where __name__ == "__main__")
 # and then calls `uvicorn.run`. the call to run() reruns the file,
@@ -155,6 +185,7 @@ def consumer():
 # server uses
 if __name__ == "server":
     initial_url_count = sqlite_helpers.get_number_of_entries(DATABASE_FILE)
+    MetricsHandler.init()
     MetricsHandler.url_count.inc(initial_url_count)
     consumer_thread = Thread(target=consumer, daemon=True)
     consumer_thread.start()