From d97a19e97d5418d9bcb2086addf0d94a66089b52 Mon Sep 17 00:00:00 2001
From: Ruslan Kuprieiev <kupruser@gmail.com>
Date: Mon, 9 Aug 2021 13:53:38 +0300
Subject: [PATCH] refodb: cache ref objects (#6392)

* refodb: cache ref objects

* odb: ignore FileNotFoundError when removing
---
 dvc/objects/db/base.py      | 4 +++-
 dvc/objects/db/reference.py | 7 +++++++
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/dvc/objects/db/base.py b/dvc/objects/db/base.py
index 3f1e3face..ef3993393 100644
--- a/dvc/objects/db/base.py
+++ b/dvc/objects/db/base.py
@@ -2,6 +2,7 @@ import itertools
 import logging
 import os
 from concurrent.futures import ThreadPoolExecutor
+from contextlib import suppress
 from copy import copy
 from typing import TYPE_CHECKING, Optional
 
@@ -203,7 +204,8 @@ class ObjectDB:
             obj.check(self, check_hash=check_hash)
         except ObjectFormatError:
             logger.warning("corrupted cache file '%s'.", obj.path_info)
-            self.fs.remove(obj.path_info)
+            with suppress(FileNotFoundError):
+                self.fs.remove(obj.path_info)
             raise
 
         if check_hash:
diff --git a/dvc/objects/db/reference.py b/dvc/objects/db/reference.py
index 12e46b8ba..0c1654ac0 100644
--- a/dvc/objects/db/reference.py
+++ b/dvc/objects/db/reference.py
@@ -27,10 +27,15 @@ class ReferenceObjectDB(ObjectDB):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self._fs_cache: Dict[tuple, "BaseFileSystem"] = {}
+        self._obj_cache: Dict["HashInfo", "ReferenceHashFile"] = {}
 
     def get(self, hash_info: "HashInfo"):
         if hash_info.isdir:
             return super().get(hash_info)
+        try:
+            return self._obj_cache[hash_info]
+        except KeyError:
+            pass
         path_info = self.hash_to_path(hash_info.value)
         try:
             with self.fs.open(path_info, "rb") as fobj:
@@ -44,6 +49,7 @@ class ReferenceObjectDB(ObjectDB):
         except ObjectFormatError:
             self.fs.remove(path_info)
             raise
+        self._obj_cache[hash_info] = ref_file
         return ref_file
 
     def _add_file(
@@ -60,6 +66,7 @@ class ReferenceObjectDB(ObjectDB):
                 from_fs, from_info, to_info, hash_info, move
             )
         ref_file = ReferenceHashFile(from_info, from_fs, hash_info)
+        self._obj_cache[hash_info] = ref_file
         ref_fobj = io.BytesIO(ref_file.to_bytes())
         ref_fobj.seek(0)
         try: