From dc1d3e880aa2f0779e8ad18de6afc5543f5a3947 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Saugat=20Pachhai=20=28=E0=A4=B8=E0=A5=8C=E0=A4=97=E0=A4=BE?=
 =?UTF-8?q?=E0=A4=A4=29?= <suagatchhetri@outlook.com>
Date: Fri, 13 May 2022 10:20:14 +0545
Subject: [PATCH] dvc.objects.cache: set type in a more compatible way

It turns out that diskcache.Disk saves passed arguments to the database
and uses that all the time. So this means that we are not forward
compatible as it breaks moving from newer version of DVC to an older
one. We don't promise this, but this breaks our benchmarks which
are probably not isolated enough.

Regarding the use of cache, it is a bit iffy to extend it this way.
We only need this type for error messages. We could consider removing
this, but since the error message is documented and it's not too hard
to extend this :), I decided to preserve the message.
---
 dvc/data/db/index.py             |  4 +---
 dvc/objects/cache.py             | 10 +++++++---
 tests/unit/objects/test_cache.py |  2 +-
 3 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/dvc/data/db/index.py b/dvc/data/db/index.py
index 60c095c8c..21bb3b12f 100644
--- a/dvc/data/db/index.py
+++ b/dvc/data/db/index.py
@@ -95,9 +95,7 @@ class ObjectDBIndex(ObjectDBIndexBase):
         self.index_dir = os.path.join(tmp_dir, self.INDEX_DIR, name)
         makedirs(self.index_dir, exist_ok=True)
         self.fs = LocalFileSystem()
-        cache = Cache(
-            self.index_dir, eviction_policy="none", disk_type="index"
-        )
+        cache = Cache(self.index_dir, eviction_policy="none", type="index")
         self.index = Index.fromcache(cache)
 
     def __iter__(self):
diff --git a/dvc/objects/cache.py b/dvc/objects/cache.py
index 0aa92d2c4..5f61535ca 100644
--- a/dvc/objects/cache.py
+++ b/dvc/objects/cache.py
@@ -35,9 +35,8 @@ def translate_pickle_error(fn):
 class Disk(disk):
     """Reraise pickle-related errors as DiskError."""
 
-    def __init__(self, *args, type=None, **kwargs):
-        super().__init__(*args, **kwargs)
-        self._type = type or os.path.basename(self._directory)
+    # we need type to differentiate cache for better error messages
+    _type: str
 
     put = translate_pickle_error(disk.put)
     get = translate_pickle_error(disk.get)
@@ -53,9 +52,14 @@ class Cache(diskcache.Cache):
         directory: str = None,
         timeout: int = 60,
         disk: disk = Disk,  # pylint: disable=redefined-outer-name
+        type: str = None,
         **settings: Any,
     ) -> None:
         settings.setdefault("disk_pickle_protocol", 4)
         super().__init__(
             directory=directory, timeout=timeout, disk=disk, **settings
         )
+        self.disk._type = self._type = type or os.path.basename(self.directory)
+
+    def __getstate__(self):
+        return (*super().__getstate__(), self._type)
diff --git a/tests/unit/objects/test_cache.py b/tests/unit/objects/test_cache.py
index 39d59d65c..5bcb0bd7c 100644
--- a/tests/unit/objects/test_cache.py
+++ b/tests/unit/objects/test_cache.py
@@ -17,7 +17,7 @@ def test_pickle_protocol_error(tmp_dir, disk_type):
     cache = Cache(
         directory,
         disk_pickle_protocol=pickle.HIGHEST_PROTOCOL + 1,
-        disk_type=disk_type,
+        type=disk_type,
     )
     with pytest.raises(DiskError) as exc, cache as cache:
         set_value(cache, "key", ("value1", "value2"))